Update the TPU dialect binding extension to follow MLIR guidelines

The way MLIR dialects are allowed to be extended in Python has recently
changed (in llvm/llvm-project#68853), so we have
to update our bindings.

PiperOrigin-RevId: 575796552

Author: apaszke

jaxlib/mosaic/python/BUILD

@@ -48,7 +48,6 @@ py_library(
     name = "tpu_dialect",
     srcs = [
         "_tpu_gen.py",
-        "_tpu_ops_ext.py",
         "tpu.py",
     ],
     visibility = ["//visibility:public"],

jaxlib/mosaic/python/_tpu_ops_ext.py

@@ -17,6 +17,38 @@
 # flake8: noqa: F401
 # flake8: noqa: F403
 
+
 # pylint: disable=g-bad-import-order
 from ._tpu_gen import *  # pylint: disable=wildcard-import
+from ._tpu_gen import _Dialect
 from jaxlib.mlir._mlir_libs._tpu_ext import *  # pylint: disable=wildcard-import
+try:
+  from jaxlib.mlir.dialects._ods_common import _cext
+except ImportError:
+  from mlir.dialects._ods_common import _cext
+
+
+@_cext.register_operation(_Dialect, replace=True)
+class TraceOp(TraceOp):
+  """An extension to the automatically generated TraceOp bindings."""
+
+  def __init__(self, results, message, level, *, loc=None, ip=None):
+    super().__init__(results, message, level, loc=loc, ip=ip)
+    self.regions[0].blocks.append(*[])  # Append the block.
+
+  @property
+  def body(self):
+    return self.regions[0].blocks[0]
+
+
+@_cext.register_operation(_Dialect, replace=True)
+class RegionOp(RegionOp):
+  """An extension to the automatically generated RegionOp bindings."""
+
+  def __init__(self, *, loc=None, ip=None):
+    super().__init__([], loc=loc, ip=ip)
+    self.regions[0].blocks.append()  # Append the block.
+
+  @property
+  def body(self):
+    return self.regions[0].blocks[0]

jaxlib/mosaic/python/tpu.py

@@ -254,7 +254,6 @@ def prepare_wheel(sources_path: pathlib.Path, *, cpu, include_gpu_plugin_extensi
           "__main__/jaxlib/mosaic/python/apply_vector_layout.py",
           "__main__/jaxlib/mosaic/python/infer_memref_layout.py",
           "__main__/jaxlib/mosaic/python/tpu.py",
-          "__main__/jaxlib/mosaic/python/_tpu_ops_ext.py",
       ],
   )
   # TODO (sharadmv,skyewm): can we avoid patching this file?