Reverts f44bc6e

PiperOrigin-RevId: 819903461

Author: jakeharmon8

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

@@ -56,10 +56,8 @@
 # TODO(joelwee): Remove this once jaxlib 0.8 is the minimum.
 try:
   from jaxlib.mlir.dialects import mpmd
-  from jaxlib.mlir.dialects import ub
 except ImportError:
   mpmd: Any = None  # type: ignore[no-redef]
-  ub: Any = None  # type: ignore[no-redef]
 from jaxlib.mlir.dialects import sdy
 
 # Alias that is set up to abstract away the transition from MHLO to StableHLO.

jaxlib/BUILD

@@ -86,7 +86,6 @@ pytype_strict_library(
         "//jaxlib/mlir:sdy_dialect",
         "//jaxlib/mlir:sparse_tensor_dialect",
         "//jaxlib/mlir:stablehlo_dialect",
-        "//jaxlib/mlir:ub_dialect",
         "//jaxlib/mlir:vector_dialect",
         "//jaxlib/mlir/_mlir_libs:_jax_mlir_ext",
         "//jaxlib/mosaic",

jaxlib/mlir/BUILD.bazel

@@ -87,19 +87,6 @@ symlink_inputs(
     ],
 )
 
-symlink_inputs(
-    name = "ub_dialect",
-    rule = py_library,
-    symlinked_inputs = {"srcs": {
-        "dialects": ["@llvm-project//mlir/python:UbPyFiles"],
-    }},
-    deps = [
-        ":core",
-        ":ir",
-        ":mlir",
-    ],
-)
-
 symlink_inputs(
     name = "math_dialect",
     rule = py_library,

jaxlib/tools/build_wheel.py

@@ -311,7 +311,6 @@ def prepare_wheel(wheel_sources_path: pathlib.Path, *, cpu, wheel_sources):
           f"{source_file_prefix}jaxlib/mlir/dialects/sdy.py",
           f"{source_file_prefix}jaxlib/mlir/dialects/sparse_tensor.py",
           f"{source_file_prefix}jaxlib/mlir/dialects/stablehlo.py",
-          f"{source_file_prefix}jaxlib/mlir/dialects/ub.py",
           f"{source_file_prefix}jaxlib/mlir/dialects/vector.py",
           f"{source_file_prefix}jaxlib/mlir/dialects/nvgpu.py",
           f"{source_file_prefix}jaxlib/mlir/dialects/nvvm.py",

tests/mosaic/gpu_dialect_test.py

@@ -24,12 +24,12 @@
 from jax._src.interpreters import mlir as mlir_interpreter
 from jax._src.lib.mlir import ir
 from jax._src.lib.mlir.dialects import arith
+from jax._src.lib.mlir.dialects import builtin
 from jax._src.lib.mlir.dialects import gpu
 from jax._src.lib.mlir.dialects import llvm
 from jax._src.lib.mlir.dialects import memref
 from jax._src.lib.mlir.dialects import nvvm
 from jax._src.lib.mlir.dialects import scf
-from jax._src.lib.mlir.dialects import ub
 from jax._src.lib.mlir.dialects import vector
 from jax.experimental.mosaic import gpu as mgpu
 from jax.experimental.mosaic.gpu import dialect_lowering as lowering
@@ -86,7 +86,8 @@ def workgroup_ptr_ty() -> ir.Type:
 
 def undefs(*tys: ir.Type) -> list[ir.Value]:
   """Returns a list of undefined values of the given types."""
-  return [ub.poison(ty) for ty in tys]
+  # TODO(allanrenucci): Use `ub.poison` once Python bindings are available.
+  return [builtin.unrealized_conversion_cast([ty], []) for ty in tys]
 
 
 class MosaicGpuTest(parameterized.TestCase):