From 60cb976e82f30bb98b0042e837b7c52a7ea3052e Mon Sep 17 00:00:00 2001
From: Baruch Tabanpour <btaba@google.com>
Date: Thu, 18 Sep 2025 09:51:50 -0700
Subject: [PATCH] Add assertion for leading dims in MJX-Warp.

PiperOrigin-RevId: 808624659
Change-Id: I259f9cd753bdffbe465ae807ec957f123e8d589e
---
 mjx/mujoco/mjx/warp/ffi.py          | 60 ++++++++++++++++++++++++++---
 mjx/mujoco/mjx/warp/forward_test.py | 28 ++++++++++++++
 2 files changed, 82 insertions(+), 6 deletions(-)

diff --git a/mjx/mujoco/mjx/warp/ffi.py b/mjx/mujoco/mjx/warp/ffi.py
index f406c408..9f84a75a 100644
--- a/mjx/mujoco/mjx/warp/ffi.py
+++ b/mjx/mujoco/mjx/warp/ffi.py
@@ -183,11 +183,8 @@ def format_args_for_warp(func, verbose=False):
   return wrapper
 
 
-def _get_mapping_from_tree_path(
-    path: jax.tree_util.KeyPath,
-    mapping: dict[str, int],
-) -> Optional[int]:
-  """Gets the mapped value from a tree path."""
+def _tree_path_to_attr_str(path: jax.tree_util.KeyPath) -> str:
+  """Converts a tree path to a dataclass attribute string."""
   if not isinstance(path, tuple):
     raise NotImplementedError(
         f'Parsing for jax tree path {path} not implemented.'
@@ -200,8 +197,15 @@ def _get_mapping_from_tree_path(
 
   assert all(isinstance(p, jax.tree_util.GetAttrKey) for p in path)
   path = [p for p in path if p.name != '_impl']
-  attr = '__'.join(p.name for p in path)
+  return '__'.join(p.name for p in path)
 
+
+def _get_mapping_from_tree_path(
+    path: jax.tree_util.KeyPath,
+    mapping: dict[str, int],
+) -> Optional[int]:
+  """Gets the mapped value from a tree path."""
+  attr = _tree_path_to_attr_str(path)
   # None if the MJX public field is not present in the MJX-Warp mapping.
   return mapping.get(attr)
 
@@ -302,6 +306,43 @@ def _maybe_broadcast_to(
   return leaf
 
 
+def _check_leading_dim(
+    path: jax.tree_util.KeyPath,
+    leaf: Any,
+    expected_batch_dim: int,
+    expected_nconmax: int,
+    expected_njmax: int,
+):
+  """Asserts that the batch dimension of a leaf node matches the expected batch dimension."""
+  has_batch_dim = _get_mapping_from_tree_path(
+      path, mjx_warp_types._BATCH_DIM['Data']
+  )
+  attr = _tree_path_to_attr_str(path)
+  if has_batch_dim and leaf.shape[0] != expected_batch_dim:
+    raise ValueError(
+        f'Leaf node batch size ({leaf.shape[0]}) and expected batch size'
+        f' ({expected_batch_dim}) do not match for field {attr}.'
+    )
+  if (
+      not has_batch_dim
+      and attr.startswith('contact__')
+      and leaf.shape[0] != expected_nconmax
+  ):
+    raise ValueError(
+        f'Leaf node leading dim ({leaf.shape[0]}) does not match nconmax'
+        f' ({expected_nconmax}) for field {attr}.'
+    )
+  if (
+      not has_batch_dim
+      and attr.startswith('efc__')
+      and leaf.shape[0] != expected_njmax
+  ):
+    raise ValueError(
+        f'Leaf node leading dim ({leaf.shape[0]}) does not match njmax'
+        f' ({expected_njmax}) for field {attr}.'
+    )
+
+
 def marshal_custom_vmap(vmap_func):
   """Marshal fields for a custom vmap into an MuJoCo Warp function."""
 
@@ -316,6 +357,13 @@ def marshal_custom_vmap(vmap_func):
         ),
         d, is_batched[1],  # fmt: skip
     )
+    # Check leading dimensions.
+    jax.tree.map_with_path(
+        lambda path, x: _check_leading_dim(
+            path, x, d_broadcast.qpos.shape[0], d._impl.nconmax, d._impl.njmax  # pylint: disable=protected-access
+        ),
+        d_broadcast,
+    )
     # Flatten batch dims into the first axis if the vmap was nested.
     m_flat = jax.tree.map_with_path(
         lambda path, x: _flatten_batch_dim(
diff --git a/mjx/mujoco/mjx/warp/forward_test.py b/mjx/mujoco/mjx/warp/forward_test.py
index 6f50f22d..ea95bfdb 100644
--- a/mjx/mujoco/mjx/warp/forward_test.py
+++ b/mjx/mujoco/mjx/warp/forward_test.py
@@ -261,6 +261,34 @@ class StepTest(parameterized.TestCase):
       tu.assert_attr_eq(dx, d, 'mocap_quat')
       tu.assert_attr_eq(dx, d, 'sensordata')
 
+  def test_step_leading_dim_mismatch(self):
+    if not _FORCE_TEST:
+      if not mjxw.WARP_INSTALLED:
+        self.skipTest('Warp not installed.')
+      if not io.has_cuda_gpu_device():
+        self.skipTest('No CUDA GPU device available.')
+
+    xml = 'humanoid/humanoid.xml'
+    batch_size = 7
+
+    m = test_util.load_test_file(xml)
+    mx = mjx.put_model(m, impl='warp')
+
+    worldids = jp.arange(batch_size)
+    dx_batch = jax.vmap(functools.partial(tu.make_data, m))(worldids)
+    dx_batch_orig = dx_batch
+
+    with self.assertRaises(ValueError):
+      dx_batch = dx_batch.replace(qpos=dx_batch.qpos[1:])
+      _ = jax.jit(jax.vmap(forward.step, in_axes=(None, 0)))(mx, dx_batch)
+
+    dx_batch = dx_batch_orig
+    with self.assertRaises(ValueError):
+      dx_batch = dx_batch.tree_replace(
+          {'_impl.contact__pos': dx_batch._impl.contact__pos[1:]}
+      )
+      _ = jax.jit(jax.vmap(forward.step, in_axes=(None, 0)))(mx, dx_batch)
+
 
 if __name__ == '__main__':
   absltest.main()