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/_tpu_ops_ext.py

@@ -17,6 +17,35 @@
 # flake8: noqa: F401
 # flake8: noqa: F403
 
+from mlir.dialects._ods_common import _cext
+
 # 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
+
+
+@_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]