Partially rolling forward #22998

Reverts 322d0c2

PiperOrigin-RevId: 670173462

jax/_src/lib/mlir/dialects/__init__.py

@@ -13,35 +13,49 @@
 # limitations under the License.
 
 # ruff: noqa: F401
-from typing import Any
 
-import jaxlib.mlir.dialects.arith as arith
-import jaxlib.mlir.dialects.builtin as builtin
-import jaxlib.mlir.dialects.chlo as chlo
-import jaxlib.mlir.dialects.func as func
-import jaxlib.mlir.dialects.math as math
-import jaxlib.mlir.dialects.memref as memref
-import jaxlib.mlir.dialects.mhlo as mhlo
-import jaxlib.mlir.dialects.scf as scf
+from typing import Any, TYPE_CHECKING
+
+if TYPE_CHECKING:
+  from jaxlib.mlir.dialects import arith as arith
+  from jaxlib.mlir.dialects import builtin as builtin
+  from jaxlib.mlir.dialects import chlo as chlo
+  from jaxlib.mlir.dialects import func as func
+  from jaxlib.mlir.dialects import gpu as gpu
+  from jaxlib.mlir.dialects import llvm as llvm
+  from jaxlib.mlir.dialects import math as math
+  from jaxlib.mlir.dialects import memref as memref
+  from jaxlib.mlir.dialects import mhlo as mhlo
+  from jaxlib.mlir.dialects import nvgpu as nvgpu
+  from jaxlib.mlir.dialects import nvvm as nvvm
+  from jaxlib.mlir.dialects import scf as scf
+  from jaxlib.mlir.dialects import sparse_tensor as sparse_tensor
+  from jaxlib.mlir.dialects import vector as vector
+else:
+  from jax._src import lazy_loader as _lazy
+  __getattr__, __dir__, __all__ = _lazy.attach("jaxlib.mlir.dialects", [
+      "arith",
+      "builtin",
+      "chlo",
+      "func",
+      "gpu",
+      "llvm",
+      "math",
+      "memref",
+      "mhlo",
+      "nvgpu",
+      "nvvm",
+      "scf",
+      "sparse_tensor",
+      "vector",
+  ])
+  del _lazy
+
 # TODO(bartchr): Once JAX is released with SDY, remove the try/except.
 try:
-  import jaxlib.mlir.dialects.sdy as sdy
+  from jaxlib.mlir.dialects import sdy as sdy
 except ImportError:
   sdy: Any = None  # type: ignore[no-redef]
-import jaxlib.mlir.dialects.sparse_tensor as sparse_tensor
-import jaxlib.mlir.dialects.vector as vector
-try:
-  # pytype: disable=import-error
-  import jaxlib.mlir.dialects.gpu as gpu
-  import jaxlib.mlir.dialects.nvgpu as nvgpu
-  import jaxlib.mlir.dialects.nvvm as nvvm
-  import jaxlib.mlir.dialects.llvm as llvm
-  # pytype: enable=import-error
-except ImportError:
-  pass
-
-from jax._src import lib
-
 
 # Alias that is set up to abstract away the transition from MHLO to StableHLO.
-import jaxlib.mlir.dialects.stablehlo as hlo
+from jaxlib.mlir.dialects import stablehlo as hlo