Update kepler_jax.py (#9)

Looks like they moved jaxlib.mhlo_helpers to jaxlib.hlo_helpers jax-ml/jax@b8ae8e3#diff-50658d597212b4ce070b8bd8c1fc522deeee1845ba387a0a5b507c446e8ea12a
dfm · Mar 21, 2023 · 1cb4c39 · 1cb4c39
1 parent 4796fee
commit 1cb4c39
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/kepler_jax/kepler_jax.py b/src/kepler_jax/kepler_jax.py
@@ -10,7 +10,7 @@
 from jax.abstract_arrays import ShapedArray
 from jax.interpreters import ad, batching, mlir, xla
 from jax.lib import xla_client
-from jaxlib.mhlo_helpers import custom_call
+from jaxlib.hlo_helpers import custom_call
 
 # Register the CPU XLA custom calls
 from . import cpu_ops