Fix misuses of jax._src.lib.mlir.

Add mpmd dialect to jax.extend.mlir.dialects.

PiperOrigin-RevId: 826263298

Author: hawkinsp

jax_tpu_embedding/sparsecore/lib/core/primitives/BUILD

@@ -32,6 +32,9 @@ pytype_strict_library(
         pypi_requirement("jax"),
         pypi_requirement("jax/_src/lib"),
         pypi_requirement("jax/extend"),
+        pypi_requirement("jax/extend/mlir"),
+        pypi_requirement("jax/extend/mlir:ir"),
+        pypi_requirement("jax/extend/mlir/dialects:stablehlo_dialect"),
         pypi_requirement("numpy"),
     ],
 )
@@ -43,6 +46,9 @@ pytype_strict_library(
         pypi_requirement("jax"),
         pypi_requirement("jax/_src/lib"),
         pypi_requirement("jax/extend"),
+        pypi_requirement("jax/extend/mlir"),
+        pypi_requirement("jax/extend/mlir:ir"),
+        pypi_requirement("jax/extend/mlir/dialects:stablehlo_dialect"),
         pypi_requirement("numpy"),
     ],
 )
@@ -58,6 +64,10 @@ pytype_strict_library(
         pypi_requirement("jax"),
         pypi_requirement("jax/_src/lib"),
         pypi_requirement("jax/extend"),
+        pypi_requirement("jax/extend/mlir"),
+        pypi_requirement("jax/extend/mlir:ir"),
+        pypi_requirement("jax/extend/mlir/dialects:func_dialect"),
+        pypi_requirement("jax/extend/mlir/dialects:stablehlo_dialect"),
         pypi_requirement("numpy"),
     ],
 )
@@ -73,6 +83,10 @@ pytype_strict_library(
         pypi_requirement("jax"),
         pypi_requirement("jax/_src/lib"),
         pypi_requirement("jax/extend"),
+        pypi_requirement("jax/extend/mlir"),
+        pypi_requirement("jax/extend/mlir:ir"),
+        pypi_requirement("jax/extend/mlir/dialects:func_dialect"),
+        pypi_requirement("jax/extend/mlir/dialects:stablehlo_dialect"),
         pypi_requirement("numpy"),
     ],
 )
@@ -88,6 +102,10 @@ pytype_strict_library(
         pypi_requirement("jax"),
         pypi_requirement("jax/_src/lib"),
         pypi_requirement("jax/extend"),
+        pypi_requirement("jax/extend/mlir"),
+        pypi_requirement("jax/extend/mlir:ir"),
+        pypi_requirement("jax/extend/mlir/dialects:func_dialect"),
+        pypi_requirement("jax/extend/mlir/dialects:stablehlo_dialect"),
         pypi_requirement("numpy"),
     ],
 )
@@ -101,6 +119,10 @@ pytype_strict_library(
         pypi_requirement("jax"),
         pypi_requirement("jax/_src/lib"),
         pypi_requirement("jax/extend"),
+        pypi_requirement("jax/extend/mlir"),
+        pypi_requirement("jax/extend/mlir:ir"),
+        pypi_requirement("jax/extend/mlir/dialects:func_dialect"),
+        pypi_requirement("jax/extend/mlir/dialects:stablehlo_dialect"),
         pypi_requirement("numpy"),
     ],
 )
@@ -116,6 +138,10 @@ pytype_strict_library(
         pypi_requirement("jax"),
         pypi_requirement("jax/_src/lib"),
         pypi_requirement("jax/extend"),
+        pypi_requirement("jax/extend/mlir"),
+        pypi_requirement("jax/extend/mlir:ir"),
+        pypi_requirement("jax/extend/mlir/dialects:func_dialect"),
+        pypi_requirement("jax/extend/mlir/dialects:stablehlo_dialect"),
         pypi_requirement("numpy"),
     ],
 )
@@ -131,6 +157,10 @@ pytype_strict_library(
         pypi_requirement("jax"),
         pypi_requirement("jax/_src/lib"),
         pypi_requirement("jax/extend"),
+        pypi_requirement("jax/extend/mlir"),
+        pypi_requirement("jax/extend/mlir:ir"),
+        pypi_requirement("jax/extend/mlir/dialects:func_dialect"),
+        pypi_requirement("jax/extend/mlir/dialects:stablehlo_dialect"),
         pypi_requirement("numpy"),
     ],
 )
@@ -143,6 +173,11 @@ pytype_strict_library(
     deps = [
         pypi_requirement("jax"),
         pypi_requirement("jax/_src/lib"),
+        pypi_requirement("jax/extend"),
+        pypi_requirement("jax/extend/mlir"),
+        pypi_requirement("jax/extend/mlir:ir"),
+        pypi_requirement("jax/extend/mlir/dialects:func_dialect"),
+        pypi_requirement("jax/extend/mlir/dialects:stablehlo_dialect"),
     ],
 )
 
@@ -156,6 +191,9 @@ pytype_strict_library(
         pypi_requirement("jax"),
         pypi_requirement("jax/_src/lib"),
         pypi_requirement("jax/extend"),
+        pypi_requirement("jax/extend/mlir"),
+        pypi_requirement("jax/extend/mlir:ir"),
+        pypi_requirement("jax/extend/mlir/dialects:stablehlo_dialect"),
         pypi_requirement("numpy"),
     ],
 )

jax_tpu_embedding/sparsecore/lib/core/primitives/local_sparse_dense_matmul.py

@@ -18,9 +18,9 @@
 
 import jax
 from jax import core
-from jax._src.lib.mlir import ir
-from jax._src.lib.mlir.dialects import hlo
 import jax.extend as jex
+from jax.extend.mlir import ir
+from jax.extend.mlir.dialects import stablehlo as hlo
 from jax.interpreters import mlir
 from jax.interpreters import xla
 import jax.numpy as jnp

jax_tpu_embedding/sparsecore/lib/core/primitives/optimizers_computation.py

@@ -13,9 +13,9 @@
 # limitations under the License.
 """Defines common optimizers for embedding lookups."""
 
-from jax._src.lib.mlir import ir
-from jax._src.lib.mlir.dialects import func as func_dialect
-from jax._src.lib.mlir.dialects import hlo
+from jax.extend.mlir import ir
+from jax.extend.mlir.dialects import func as func_dialect
+from jax.extend.mlir.dialects import stablehlo as hlo
 from jax.interpreters import mlir
 
 

jax_tpu_embedding/sparsecore/lib/core/primitives/sparse_dense_matmul_csr.py

@@ -18,9 +18,9 @@
 
 import jax
 from jax import core
-from jax._src.lib.mlir import ir
-from jax._src.lib.mlir.dialects import hlo
 import jax.extend as jex
+from jax.extend.mlir import ir
+from jax.extend.mlir.dialects import stablehlo as hlo
 from jax.interpreters import mlir
 from jax.interpreters import xla
 import jax.numpy as jnp

jax_tpu_embedding/sparsecore/lib/core/primitives/sparse_dense_matmul_grad_with_adagrad.py

@@ -25,10 +25,10 @@
 from typing import Tuple
 
 import jax
-from jax._src.lib.mlir import ir
-from jax._src.lib.mlir.dialects import func as func_dialect
-from jax._src.lib.mlir.dialects import hlo
 import jax.extend as jex
+from jax.extend.mlir import ir
+from jax.extend.mlir.dialects import func as func_dialect
+from jax.extend.mlir.dialects import stablehlo as hlo
 from jax.interpreters import mlir
 from jax.interpreters import xla
 from jax_tpu_embedding.sparsecore.lib.core import constants

jax_tpu_embedding/sparsecore/lib/core/primitives/sparse_dense_matmul_grad_with_adagrad_momentum.py

@@ -27,10 +27,10 @@
 from typing import Tuple
 
 import jax
-from jax._src.lib.mlir import ir
-from jax._src.lib.mlir.dialects import func as func_dialect
-from jax._src.lib.mlir.dialects import hlo
 import jax.extend as jex
+from jax.extend.mlir import ir
+from jax.extend.mlir.dialects import func as func_dialect
+from jax.extend.mlir.dialects import stablehlo as hlo
 from jax.interpreters import mlir
 from jax.interpreters import xla
 from jax_tpu_embedding.sparsecore.lib.core import constants

jax_tpu_embedding/sparsecore/lib/core/primitives/sparse_dense_matmul_grad_with_adam.py

@@ -27,10 +27,10 @@
 from typing import Tuple
 
 import jax
-from jax._src.lib.mlir import ir
-from jax._src.lib.mlir.dialects import func as func_dialect
-from jax._src.lib.mlir.dialects import hlo
 import jax.extend as jex
+from jax.extend.mlir import ir
+from jax.extend.mlir.dialects import func as func_dialect
+from jax.extend.mlir.dialects import stablehlo as hlo
 from jax.interpreters import mlir
 from jax.interpreters import xla
 from jax_tpu_embedding.sparsecore.lib.core import constants

jax_tpu_embedding/sparsecore/lib/core/primitives/sparse_dense_matmul_grad_with_ftrl.py

@@ -28,10 +28,10 @@
 from typing import Tuple
 
 import jax
-from jax._src.lib.mlir import ir
-from jax._src.lib.mlir.dialects import func as func_dialect
-from jax._src.lib.mlir.dialects import hlo
 import jax.extend as jex
+from jax.extend.mlir import ir
+from jax.extend.mlir.dialects import func as func_dialect
+from jax.extend.mlir.dialects import stablehlo as hlo
 from jax.interpreters import mlir
 from jax.interpreters import xla
 from jax_tpu_embedding.sparsecore.lib.core import constants

jax_tpu_embedding/sparsecore/lib/core/primitives/sparse_dense_matmul_grad_with_laprop.py

@@ -26,10 +26,10 @@
 from typing import Tuple
 
 import jax
-from jax._src.lib.mlir import ir
-from jax._src.lib.mlir.dialects import func as func_dialect
-from jax._src.lib.mlir.dialects import hlo
 import jax.extend as jex
+from jax.extend.mlir import ir
+from jax.extend.mlir.dialects import func as func_dialect
+from jax.extend.mlir.dialects import stablehlo as hlo
 from jax.interpreters import mlir
 from jax.interpreters import xla
 from jax_tpu_embedding.sparsecore.lib.core import constants

jax_tpu_embedding/sparsecore/lib/core/primitives/sparse_dense_matmul_grad_with_sgd.py

@@ -17,10 +17,10 @@
 import json
 
 import jax
-from jax._src.lib.mlir import ir
-from jax._src.lib.mlir.dialects import func as func_dialect
-from jax._src.lib.mlir.dialects import hlo
 import jax.extend as jex
+from jax.extend.mlir import ir
+from jax.extend.mlir.dialects import func as func_dialect
+from jax.extend.mlir.dialects import stablehlo as hlo
 from jax.interpreters import mlir
 from jax.interpreters import xla
 from jax_tpu_embedding.sparsecore.lib.core import constants