From e5a236774d38e9ae79100673d3a48537bed14449 Mon Sep 17 00:00:00 2001
From: Baruch Tabanpour <btaba@google.com>
Date: Fri, 20 Feb 2026 09:20:38 -0800
Subject: [PATCH] Add has_side_effect for pmap bvh refit in mjx-warp.

PiperOrigin-RevId: 872934001
Change-Id: If48cab439718b9e783c4b1da4fb1b1a08521a602
---
 .../mjx/third_party/warp/_src/jax_experimental/ffi.py       | 6 +++++-
 mjx/mujoco/mjx/warp/bvh.py                                  | 3 ++-
 mjx/mujoco/mjx/warp/collision_driver.py                     | 1 +
 mjx/mujoco/mjx/warp/ffi.py                                  | 2 ++
 mjx/mujoco/mjx/warp/forward.py                              | 2 ++
 mjx/mujoco/mjx/warp/render.py                               | 1 +
 mjx/mujoco/mjx/warp/smooth.py                               | 2 ++
 7 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py b/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py
index 43f63397..9e2e58af 100644
--- a/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py
+++ b/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py
@@ -434,6 +434,7 @@ class FfiCallable:
         stage_out_argnames,
         graph_cache_max,
         module_preload_mode,
+        has_side_effect=False,
     ):
         self.func = func
         self.name = generate_unique_name(func)
@@ -442,6 +443,7 @@ class FfiCallable:
         self.graph_mode = graph_mode
         self.output_dims = output_dims
         self.module_preload_mode = module_preload_mode
+        self.has_side_effect = has_side_effect
         self.first_array_arg = None
         self.call_id = 0
         self.call_descriptors = {}
@@ -613,7 +615,7 @@ class FfiCallable:
             out_types,
             vmap_method=vmap_method,
             input_output_aliases=self.input_output_aliases,
-            # has_side_effect=True,  # force this function to execute even if outputs aren't used
+            has_side_effect=self.has_side_effect,
         )
 
         # preload on the specified devices
@@ -1379,6 +1381,7 @@ def jax_callable(
     stage_out_argnames=None,
     graph_cache_max: int | None = None,
     module_preload_mode: ModulePreloadMode = ModulePreloadMode.CURRENT_DEVICE,
+    has_side_effect: bool = False,
 ):
     """Create a JAX callback from an annotated Python function.
 
@@ -1449,6 +1452,7 @@ def jax_callable(
                 stage_out_argnames,
                 graph_cache_max,
                 module_preload_mode,
+                has_side_effect,
             )
             _FFI_CALLABLE_REGISTRY[key] = callable
         else:
diff --git a/mjx/mujoco/mjx/warp/bvh.py b/mjx/mujoco/mjx/warp/bvh.py
index a86bad63..529145ee 100644
--- a/mjx/mujoco/mjx/warp/bvh.py
+++ b/mjx/mujoco/mjx/warp/bvh.py
@@ -104,6 +104,7 @@ def _refit_bvh_jax_impl(m: types.Model, d: types.Data, ctx: RenderContext):
       stage_in_argnames=set(['geom_size', 'geom_xmat', 'geom_xpos']),
       stage_out_argnames=set([]),
       graph_mode=m.opt._impl.graph_mode,
+      has_side_effect=True,
   )
   out = jf(
       d.qpos.shape[0],
@@ -122,7 +123,7 @@ def _refit_bvh_jax_impl(m: types.Model, d: types.Data, ctx: RenderContext):
       d.geom_xpos,
       ctx.key,
   )
-  d = d.tree_replace({'time': d.time + out[0]})
+  d = d.tree_replace({})
   return d
 
 
diff --git a/mjx/mujoco/mjx/warp/collision_driver.py b/mjx/mujoco/mjx/warp/collision_driver.py
index aa88f118..2ae01123 100644
--- a/mjx/mujoco/mjx/warp/collision_driver.py
+++ b/mjx/mujoco/mjx/warp/collision_driver.py
@@ -289,6 +289,7 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
       ]),
       stage_out_argnames=set([]),
       graph_mode=m.opt._impl.graph_mode,
+      has_side_effect=False,
   )
   out = jf(
       d.qpos.shape[0],
diff --git a/mjx/mujoco/mjx/warp/ffi.py b/mjx/mujoco/mjx/warp/ffi.py
index 6d0b6d7e..a485cbbb 100644
--- a/mjx/mujoco/mjx/warp/ffi.py
+++ b/mjx/mujoco/mjx/warp/ffi.py
@@ -103,6 +103,7 @@ def jax_callable_variadic_tuple(
     in_out_argnames: Optional[Sequence[str]] = None,
     stage_in_argnames: Optional[Sequence[str]] = None,
     stage_out_argnames: Optional[Sequence[str]] = None,
+    has_side_effect: bool = False,
 ):
   """Wraps a JAX callable to support variadic tuples and dataclasses."""
 
@@ -134,6 +135,7 @@ def jax_callable_variadic_tuple(
         in_out_argnames=in_out_argnames,
         stage_in_argnames=stage_in_argnames,
         stage_out_argnames=stage_out_argnames,
+        has_side_effect=has_side_effect,
     )
 
     flat_args, in_tree = jax.tree.flatten(args)
diff --git a/mjx/mujoco/mjx/warp/forward.py b/mjx/mujoco/mjx/warp/forward.py
index db52a8cd..e48219a8 100644
--- a/mjx/mujoco/mjx/warp/forward.py
+++ b/mjx/mujoco/mjx/warp/forward.py
@@ -1279,6 +1279,7 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
           'xquat',
       ]),
       graph_mode=m.opt._impl.graph_mode,
+      has_side_effect=False,
   )
   out = jf(
       d.qpos.shape[0],
@@ -3070,6 +3071,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'xquat',
       ]),
       graph_mode=m.opt._impl.graph_mode,
+      has_side_effect=False,
   )
   out = jf(
       d.qpos.shape[0],
diff --git a/mjx/mujoco/mjx/warp/render.py b/mjx/mujoco/mjx/warp/render.py
index 73ced206..3b2d3aee 100644
--- a/mjx/mujoco/mjx/warp/render.py
+++ b/mjx/mujoco/mjx/warp/render.py
@@ -145,6 +145,7 @@ def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContext):
       ]),
       stage_out_argnames=set([]),
       graph_mode=m.opt._impl.graph_mode,
+      has_side_effect=False,
   )
   out = jf(
       d.qpos.shape[0],
diff --git a/mjx/mujoco/mjx/warp/smooth.py b/mjx/mujoco/mjx/warp/smooth.py
index 6e1c0ea8..0e1ec813 100644
--- a/mjx/mujoco/mjx/warp/smooth.py
+++ b/mjx/mujoco/mjx/warp/smooth.py
@@ -214,6 +214,7 @@ def _kinematics_jax_impl(m: types.Model, d: types.Data):
           'xquat',
       ]),
       graph_mode=m.opt._impl.graph_mode,
+      has_side_effect=False,
   )
   out = jf(
       d.qpos.shape[0],
@@ -408,6 +409,7 @@ def _tendon_jax_impl(m: types.Model, d: types.Data):
       ]),
       stage_out_argnames=set(['ten_length']),
       graph_mode=m.opt._impl.graph_mode,
+      has_side_effect=False,
   )
   out = jf(
       d.qpos.shape[0],