Add stage_in and stage_out for MJX-Warp.

PiperOrigin-RevId: 868180029
Change-Id: I0d1b17b6023d32498766b5d1d9e5624bd31cda3d

Author: btaba

mjx/mujoco/mjx/warp/collision_driver.py

@@ -233,8 +233,6 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
       'collision_pair': d._impl.collision_pair.shape,
       'collision_pairid': d._impl.collision_pairid.shape,
       'collision_worldid': d._impl.collision_worldid.shape,
-      'geom_xmat': d.geom_xmat.shape,
-      'geom_xpos': d.geom_xpos.shape,
       'nacon': d._impl.nacon.shape,
       'ncollision': d._impl.ncollision.shape,
       'contact__dim': d._impl.contact__dim.shape,
@@ -253,15 +251,13 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
   }
   jf = ffi.jax_callable_variadic_tuple(
       _collision_shim,
-      num_outputs=20,
+      num_outputs=18,
       output_dims=output_dims,
       vmap_method=None,
       in_out_argnames={
           'collision_pair',
           'collision_pairid',
           'collision_worldid',
-          'geom_xmat',
-          'geom_xpos',
           'nacon',
           'ncollision',
           'contact__dim',
@@ -278,6 +274,27 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
           'contact__type',
           'contact__worldid',
       },
+      stage_in_argnames={
+          'geom_aabb',
+          'geom_friction',
+          'geom_gap',
+          'geom_margin',
+          'geom_rbound',
+          'geom_size',
+          'geom_solimp',
+          'geom_solmix',
+          'geom_solref',
+          'geom_xmat',
+          'geom_xpos',
+          'hfield_data',
+          'pair_friction',
+          'pair_gap',
+          'pair_margin',
+          'pair_solimp',
+          'pair_solref',
+          'pair_solreffriction',
+      },
+      stage_out_argnames={},
       graph_mode=m.opt._impl.graph_mode,
   )
   out = jf(
@@ -373,23 +390,21 @@ def _collision_jax_impl(m: types.Model, d: types.Data):
       '_impl.collision_pair': out[0],
       '_impl.collision_pairid': out[1],
       '_impl.collision_worldid': out[2],
-      'geom_xmat': out[3],
-      'geom_xpos': out[4],
-      '_impl.nacon': out[5],
-      '_impl.ncollision': out[6],
-      '_impl.contact__dim': out[7],
-      '_impl.contact__dist': out[8],
-      '_impl.contact__frame': out[9],
-      '_impl.contact__friction': out[10],
-      '_impl.contact__geom': out[11],
-      '_impl.contact__geomcollisionid': out[12],
-      '_impl.contact__includemargin': out[13],
-      '_impl.contact__pos': out[14],
-      '_impl.contact__solimp': out[15],
-      '_impl.contact__solref': out[16],
-      '_impl.contact__solreffriction': out[17],
-      '_impl.contact__type': out[18],
-      '_impl.contact__worldid': out[19],
+      '_impl.nacon': out[3],
+      '_impl.ncollision': out[4],
+      '_impl.contact__dim': out[5],
+      '_impl.contact__dist': out[6],
+      '_impl.contact__frame': out[7],
+      '_impl.contact__friction': out[8],
+      '_impl.contact__geom': out[9],
+      '_impl.contact__geomcollisionid': out[10],
+      '_impl.contact__includemargin': out[11],
+      '_impl.contact__pos': out[12],
+      '_impl.contact__solimp': out[13],
+      '_impl.contact__solref': out[14],
+      '_impl.contact__solreffriction': out[15],
+      '_impl.contact__type': out[16],
+      '_impl.contact__worldid': out[17],
   })
   return d
 

mjx/mujoco/mjx/warp/ffi.py

@@ -101,6 +101,8 @@ def jax_callable_variadic_tuple(
     vmap_method: Optional[str] = None,
     output_dims: Optional[dict[str, tuple[int, ...]]] = None,
     in_out_argnames: Optional[Sequence[str]] = None,
+    stage_in_argnames: Optional[Sequence[str]] = None,
+    stage_out_argnames: Optional[Sequence[str]] = None,
 ):
   """Wraps a JAX callable to support variadic tuples and dataclasses."""
 
@@ -130,6 +132,8 @@ def func_wrapper(*flat_args, **kwargs):
         vmap_method=vmap_method,
         output_dims=output_dims,
         in_out_argnames=in_out_argnames,
+        stage_in_argnames=stage_in_argnames,
+        stage_out_argnames=stage_out_argnames,
     )
 
     flat_args, in_tree = jax.tree.flatten(args)