Adds support for explicit dense/sparse mass matrices to MJX.

This increases performance, particularly for the Newton solver on TPU.

PiperOrigin-RevId: 600696483
Change-Id: If69bb9a2e21ba8dad6ca23f093ce7b7ceae644ff

mjx/mujoco/mjx/_src/support_test.py

@@ -27,6 +27,46 @@ import numpy as np
 
 class SupportTest(parameterized.TestCase):
 
+  def test_mul_m(self):
+    m = test_util.load_test_file('pendula.xml')
+    # first test sparse
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
+    d = mujoco.MjData(m)
+    # give the system a little kick to ensure we have non-identity rotations
+    d.qvel = np.random.random(m.nv)
+    mujoco.mj_step(m, d, 10)  # let dynamics get state significantly non-zero
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
+    vec = np.random.random(m.nv)
+    mjx_vec = jax.jit(mjx.mul_m)(mx, dx, jp.array(vec))
+    mj_vec = np.zeros(m.nv)
+    mujoco.mj_mulM(m, d, mj_vec, vec)
+    np.testing.assert_allclose(mjx_vec, mj_vec, atol=5e-5, rtol=5e-5)
+
+    # also check dense
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_DENSE
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
+    mjx_vec = jax.jit(mjx.mul_m)(mx, dx, jp.array(vec))
+    np.testing.assert_allclose(mjx_vec, mj_vec, atol=5e-5, rtol=5e-5)
+
+  def test_full_m(self):
+    m = test_util.load_test_file('pendula.xml')
+    # for the model to be sparse to exercise MJX full_M
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
+    d = mujoco.MjData(m)
+    # give the system a little kick to ensure we have non-identity rotations
+    d.qvel = np.random.random(m.nv)
+    mujoco.mj_step(m, d, 10)  # let dynamics get state significantly non-zero
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
+    mjx_full_m = jax.jit(support.full_m)(mx, dx)
+    mj_full_m = np.zeros((m.nv, m.nv), dtype=np.float64)
+    mujoco.mj_fullM(m, mj_full_m, d.qM)
+    np.testing.assert_allclose(mjx_full_m, mj_full_m, atol=5e-5, rtol=5e-5)
+
   @parameterized.parameters('constraints.xml', 'pendula.xml')
   def test_jac(self, fname):
     np.random.seed(0)