Add has_side_effect for pmap bvh refit in mjx-warp.

PiperOrigin-RevId: 872934001
Change-Id: If48cab439718b9e783c4b1da4fb1b1a08521a602

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:

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
 
 

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],

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)

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],

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],

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],