Update xla to use mlir rather than backend-specific-translations

sail-sg · Aug 12, 2024 · b7b8a49 · b7b8a49
1 parent f411fc2
commit b7b8a49
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/envpool/python/xla_template.py b/envpool/python/xla_template.py
@@ -23,6 +23,7 @@
 from jax.core import ShapedArray
 from jax.interpreters import xla
 from jax.lib import xla_client
+from jax import interpreters
 
 
 def _shape_with_layout(
@@ -91,10 +92,10 @@ def translation(c: Any, *args: Any, platform: str = "cpu") -> Any:
   prim.multiple_results = (len(out_specs) > 1)
   prim.def_impl(partial(xla.apply_primitive, prim))
   prim.def_abstract_eval(abstract)
-  xla.backend_specific_translations["cpu"][prim] = partial(
+  interpreters.mlir["cpu"][prim] = partial(
     translation, platform="cpu"
   )
-  xla.backend_specific_translations["gpu"][prim] = partial(
+  interpreters.mlir["gpu"][prim] = partial(
     translation, platform="gpu"
   )