Fix a bug where mjx.make_data qLDiagInv has wrong size for sparse mass matrices.

PiperOrigin-RevId: 669087189
Change-Id: I39e97264a192a621e67ae5e079334fc8be38a25d

doc/changelog.rst

@@ -42,6 +42,7 @@ MJX
 - Added ``device`` parameter to ``mjx.make_data`` to bring it to parity with ``mjx.put_model`` and ``mjx.put_data``.
 - Added support for :ref:`implicitfast integration<geIntegration>` for all cases except
   :doc:`fluid drag <computation/fluid>`.
+- Fixed a bug where ``qLDiagInv`` had the wrong size for sparse mass matrices.
 
 Bug fixes
 ^^^^^^^^^

mjx/mujoco/mjx/_src/io.py

@@ -248,7 +248,7 @@ def make_data(
         'crb': (m.nbody, 10, float),
         'qM': (m.nM, float) if support.is_sparse(m) else (m.nv, m.nv, float),
         'qLD': (m.nM, float) if support.is_sparse(m) else (m.nv, m.nv, float),
-        'qLDiagInv': (m.nM, float) if support.is_sparse(m) else (0, float),
+        'qLDiagInv': (m.nv, float) if support.is_sparse(m) else (0, float),
         'qLDiagSqrtInv': (m.nv, float),
         'bvh_aabb_dyn': (m.nbvhdynamic, 6, float),
         'bvh_active': (m.nbvh, jp.uint8),

mjx/mujoco/mjx/_src/io_test.py

@@ -229,6 +229,7 @@ class DataIOTest(parameterized.TestCase):
     nq = 22
     nbody = 5
     ncon = 46
+    nm = 64
     nv = 19
     nefc = 185
 
@@ -284,6 +285,13 @@ class DataIOTest(parameterized.TestCase):
     self.assertEqual(d.qfrc_inverse.shape, (nv,))
     self.assertEqual(d.efc_force.shape, (nefc,))
 
+    # test sparse
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
+    d = mjx.make_data(m)
+    self.assertEqual(d.qM.shape, (nm,))
+    self.assertEqual(d.qLD.shape, (nm,))
+    self.assertEqual(d.qLDiagInv.shape, (nv,))
+
   def test_put_data(self):
     """Test that put_data puts the correct data for dense and sparse."""