Fix undefined symbol wp_cuda_graph_launch in Warp backend

Harshdhall01 · Harshdhall01 · commit bb76189d35e2 · 2025-09-23T12:32:48.000Z
The Warp library exports cuda_graph_launch but MJX FFI code
was looking for wp_cuda_graph_launch (with wp_ prefix).
This caused AttributeError on second call to JAX-compiled functions.

Fixes: undefined symbol error in Warp backend with JAX compilation
diff --git a/mjx/mujoco/mjx/third_party/warp/jax_experimental/ffi.py b/mjx/mujoco/mjx/third_party/warp/jax_experimental/ffi.py
@@ -575,7 +575,7 @@ def ffi_callback(self, call_frame):
                             raise RuntimeError(f"Graph creation error: {wp.context.runtime.get_error_string()}")
                         graph.graph_exec = g
 
-                    if not wp.context.runtime.core.wp_cuda_graph_launch(graph.graph_exec, cuda_stream):
+                    if not wp.context.runtime.core.cuda_graph_launch(graph.graph_exec, cuda_stream):
                         raise RuntimeError(f"Graph launch error: {wp.context.runtime.get_error_string()}")
 
                     # early out
diff --git a/test-fix.py b/test-fix.py
@@ -0,0 +1,35 @@
+#!/usr/bin/env python3
+import jax
+from mujoco import MjModel, mjx
+
+def main() -> None:
+    host_model = MjModel.from_xml_string(model_spec_string())
+    mjx_model = mjx.put_model(host_model, impl="warp")
+    initial_sim_state = mjx.make_data(host_model, impl="warp")
+
+    print("== Calling my compiled function for the first time ==")
+    print("Result:", ball_pos_after_5_steps(mjx_model, initial_sim_state))
+
+    print("== Calling it again ==")
+    print("Result:", ball_pos_after_5_steps(mjx_model, initial_sim_state))
+
+@jax.jit
+def ball_pos_after_5_steps(model: mjx.Model, initial_sim_state: mjx.Data) -> jax.Array:
+    sim_state = initial_sim_state
+    for _ in range(5):
+        sim_state = mjx.step(model, sim_state)
+    return sim_state.qpos[:3]
+
+def model_spec_string() -> str:
+    return """<mujoco>
+        <worldbody>
+            <geom type="plane" size="10 10 0.1"/>
+            <body pos="0 0 3">
+                <geom type="sphere" size="0.2" rgba="0 1 0 1"/>
+                <joint type="free"/>
+            </body>
+        </worldbody>
+    </mujoco>"""
+
+if __name__ == "__main__":
+    main()
