Update mjx-warp ffi, prepare for render.

PiperOrigin-RevId: 829027266
Change-Id: Ifca30885fb6302e8ce61f949fdc0f56674bbe9e4

mjx/mujoco/mjx/warp/ffi.py

@@ -106,8 +106,11 @@ def jax_callable_variadic_tuple(
 
   def callable_wrapper(*args, **kwargs):
     def func_wrapper(*flat_args, **kwargs):
-      unflat_args = jax.tree.unflatten(in_tree, flat_args)
-      return func(*unflat_args, **kwargs)
+      num_inputs = in_tree.num_leaves
+      flat_inputs = flat_args[:num_inputs]
+      output_buffers = flat_args[num_inputs:]
+      unflat_args = jax.tree.unflatten(in_tree, flat_inputs)
+      return func(*unflat_args, *output_buffers, **kwargs)
 
     # Provide a flattened signature for the Warp callable machinery.
     new_signature = flatten_signature(inspect.signature(func), args)
@@ -245,7 +248,7 @@ def _squeeze_dim(leaf_expanded: Any, leaf: Any) -> Any:
   return leaf_expanded
 
 
-def marshal_jax_warp_callable(func):
+def marshal_jax_warp_callable(func, raw_output: bool = False):
   """Marshal fields into a MuJoCo Warp function."""
 
   @functools.wraps(func)
@@ -265,6 +268,9 @@ def marshal_jax_warp_callable(func):
         d,
     )
     d_expanded_result = func(m_expanded, d_expanded)
+
+    if raw_output:
+      return d_expanded_result
     d_result = jax.tree.map(_squeeze_dim, d_expanded_result, d)
     return d_result
 
@@ -350,7 +356,7 @@ def _check_leading_dim(
     )
 
 
-def marshal_custom_vmap(vmap_func):
+def marshal_custom_vmap(vmap_func, raw_output: bool = False):
   """Marshal fields for a custom vmap into an MuJoCo Warp function."""
 
   @functools.wraps(vmap_func)
@@ -387,6 +393,9 @@ def marshal_custom_vmap(vmap_func):
     d_broadcast_flat_result, out_batched = vmap_func(
         axis_size, is_batched, m_flat, d_broadcast_flat
     )
+    if raw_output:
+      return d_broadcast_flat_result, out_batched
+
     # Explicitly mark MuJoCo Warp data fields as batched after vmapping is done.
     out_batched = jax.tree.map_with_path(
         # NB: if a field is not in MuJoCo Warp, we let JAX do its magic.