From 0e61a4ac5e986449b1a6ec1913265c18cc9951ed Mon Sep 17 00:00:00 2001 From: Peter Hawkins Date: Thu, 18 Sep 2025 07:17:16 -0700 Subject: [PATCH] Remove jaxlib.hlo_helpers. This module has been deprecated for many months and its replacement is jax.ffi. PiperOrigin-RevId: 808567797 Change-Id: I5b4f4eee5e65ccd86e597c5761e68c36292b6209 --- mjx/mujoco/mjx/third_party/warp/jax_experimental/custom_call.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mjx/mujoco/mjx/third_party/warp/jax_experimental/custom_call.py b/mjx/mujoco/mjx/third_party/warp/jax_experimental/custom_call.py index 97faaec5..a90f09a2 100644 --- a/mjx/mujoco/mjx/third_party/warp/jax_experimental/custom_call.py +++ b/mjx/mujoco/mjx/third_party/warp/jax_experimental/custom_call.py @@ -119,7 +119,7 @@ def _create_jax_warp_primitive(): from jax._src.interpreters import batching from jax.interpreters import mlir from jax.interpreters.mlir import ir - from jax.jaxlib.hlo_helpers import custom_call + from jax.google.hlo_helpers import custom_call global _jax_warp_p global _cc_callback