Remove accidental exports jax.interpreters.mlir.{hlo,func_dialect}.

These are available via jax.extend.mlir.dialects.

No deprecation period because jax.interpreters.mlir is not a stable API.

PiperOrigin-RevId: 744712537

Author: hawkinsp

CHANGELOG.md

@@ -33,6 +33,9 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
   * Implemented host callback handlers for CPU and GPU devices using XLA's FFI
     and removed existing CPU/GPU handlers using XLA's custom call.
   * All APIs in `jax.lib.xla_extension` are now deprecated.
+  * `jax.interpreters.mlir.hlo` and `jax.interpreters.mlir.func_dialect`,
+    which were accidental exports, have been removed. If needed, they are
+    available from `jax.extend.mlir`.
   * Several previously-deprecated APIs have been removed, including:
     * From `jax.lib.xla_client`: `FftType`, `PaddingType`, `dtype_to_etype`,
       and `shape_from_pyval`.

jax/_src/cudnn/fused_attention_stablehlo.py

@@ -28,8 +28,8 @@
 from jax._src import xla_bridge
 from jax.interpreters import mlir
 from jax.interpreters import xla
-from jax.interpreters.mlir import hlo
-from jax.interpreters.mlir import ir
+from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import hlo
 import jax.numpy as jnp
 from jax.sharding import NamedSharding, PartitionSpec
 

jax/_src/cudnn/fusion.py

@@ -16,8 +16,8 @@
 import jax
 from jax._src import core as jax_core
 from jax.interpreters import mlir
-from jax.interpreters.mlir import hlo
-from jax.interpreters.mlir import ir
+from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import hlo
 
 
 

jax/interpreters/mlir.py

@@ -43,8 +43,6 @@
   flatten_ir_values as flatten_lowering_ir_args,  # TODO(phawkins): remove me  # noqa: F401
   flatten_ir_values as flatten_ir_values,
   unflatten_ir_values_like_types as unflatten_ir_values_like_types,
-  func_dialect as func_dialect,
-  hlo as hlo,
   i32_attr as i32_attr,
   i64_attr as i64_attr,
   ir as ir,