Remove uses of deprecated xb, xc, and xe abbreviation from jax.interp…

…reters.xla PiperOrigin-RevId: 658688048
google-deepmind · Aug 2, 2024 · 54a1eba · 54a1eba
1 parent 66d8b74
commit 54a1eba
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/haiku/_src/transform.py b/haiku/_src/transform.py
@@ -338,8 +338,8 @@ def transform(f, *, apply_rng=True) -> Transformed:
 
   return without_state(transform_with_state(f))
 
-COMPILED_FN_TYPES = (jax.interpreters.xla.xe.PjitFunction,
-                     jax.interpreters.xla.xe.PmapFunction)  # pytype: disable=name-error
+COMPILED_FN_TYPES = (jax.lib.xla_extension.PjitFunction,
+                     jax.lib.xla_extension.PmapFunction)  # pytype: disable=name-error
 
 
 def check_not_jax_transformed(f):