From 494e166f289b8e069825cfb9d6dbc03b358de7a7 Mon Sep 17 00:00:00 2001
From: Erik Frey <erikfrey@google.com>
Date: Thu, 29 Aug 2024 15:39:09 -0700
Subject: [PATCH] Fix a bug where `mjx.make_data` qLDiagInv has wrong size for
 sparse mass matrices.

PiperOrigin-RevId: 669087189
Change-Id: I39e97264a192a621e67ae5e079334fc8be38a25d
---
 doc/changelog.rst              | 1 +
 mjx/mujoco/mjx/_src/io.py      | 2 +-
 mjx/mujoco/mjx/_src/io_test.py | 8 ++++++++
 3 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/doc/changelog.rst b/doc/changelog.rst
index 439c42ec..7cc4f827 100644
--- a/doc/changelog.rst
+++ b/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
 ^^^^^^^^^
diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py
index cdd7b2a5..eb059f9f 100644
--- a/mjx/mujoco/mjx/_src/io.py
+++ b/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),
diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py
index 97c61382..fe4de217 100644
--- a/mjx/mujoco/mjx/_src/io_test.py
+++ b/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."""