Add mju_sym2dense, document future breakage of mj_fullM

PiperOrigin-RevId: 910242375
Change-Id: Ibfbdef9cfb66088723499ea257da09aee0d80938

mjx/mujoco/mjx/_src/io.py

@@ -1058,7 +1058,13 @@ def _put_data_jax(
   # convert qM and qLD if jacobian is dense
   if not support.is_sparse(m):
     impl_fields['qM'] = np.zeros((m.nv, m.nv))
-    mujoco.mj_fullM(m, impl_fields['qM'], d.qM)
+    mujoco.mju_sym2dense(
+        impl_fields['qM'],
+        d.M,
+        m.M_rownnz,
+        m.M_rowadr,
+        m.M_colind,
+    )
     # TODO(erikfrey): derive L*L' from L'*D*L instead of recomputing
     try:
       impl_fields['qLD'], _ = scipy.linalg.cho_factor(impl_fields['qM'])

mjx/mujoco/mjx/_src/io_test.py

@@ -520,7 +520,7 @@ class DataIOTest(parameterized.TestCase):
     if impl == 'jax':
       # check that qM is transformed properly
       qm = np.zeros((m.nv, m.nv), dtype=np.float64)
-      mujoco.mj_fullM(m, qm, d.qM)
+      mujoco.mju_sym2dense(qm, d.M, m.M_rownnz, m.M_rowadr, m.M_colind)
       np.testing.assert_allclose(qm, mjx.full_m(mjx.put_model(m), dx))
 
     elif impl == 'cpp':
@@ -529,7 +529,7 @@ class DataIOTest(parameterized.TestCase):
       return  # cpp does not populate other fields in _impl
     elif impl == 'warp':
       qm = np.zeros((m.nv, m.nv), dtype=np.float64)
-      mujoco.mj_fullM(m, qm, d.qM)
+      mujoco.mju_sym2dense(qm, d.M, m.M_rownnz, m.M_rowadr, m.M_colind)
       np.testing.assert_allclose(dx._impl.qM, qm)
       # TODO(taylorhowell): test efc__J
       np.testing.assert_allclose(dx._impl.efc__aref[:3], d.efc_aref[:3])
@@ -596,7 +596,7 @@ class DataIOTest(parameterized.TestCase):
     dx_from_dense = mjx.put_data(m, d, impl=impl)
     if impl == 'jax':
       qm = np.zeros((m.nv, m.nv))
-      mujoco.mj_fullM(m, qm, d.qM)
+      mujoco.mju_sym2dense(qm, d.M, m.M_rownnz, m.M_rowadr, m.M_colind)
       np.testing.assert_allclose(dx_from_dense._impl.qM, qm, atol=1e-8)
 
 

mjx/mujoco/mjx/_src/smooth_test.py

@@ -438,7 +438,7 @@ class TendonTest(parameterized.TestCase):
 
     if jacobian == JacobianType.DENSE:
       qM = np.zeros((m.nv, m.nv))  # pylint: disable=invalid-name
-      mujoco.mj_fullM(m, qM, d.qM)
+      mujoco.mju_sym2dense(qM, d.M, m.M_rownnz, m.M_rowadr, m.M_colind)
     else:
       qM = d.qM  # pylint: disable=invalid-name
     _assert_eq(dx._impl.qM, qM, 'qM')

mjx/mujoco/mjx/_src/support_test.py

@@ -64,7 +64,7 @@ class SupportTest(parameterized.TestCase):
     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)
+    mujoco.mju_sym2dense(mj_full_m, d.M, m.M_rownnz, m.M_rowadr, m.M_colind)
     np.testing.assert_allclose(mjx_full_m, mj_full_m, atol=5e-5, rtol=5e-5)
 
   @parameterized.parameters('constraints.xml', 'pendula.xml')

mjx/mujoco/mjx/warp/forward_test.py

@@ -175,7 +175,7 @@ class ForwardTest(parameterized.TestCase):
       tu.assert_attr_eq(dx._impl, d, 'crb')
 
       qm = np.zeros((m.nv, m.nv))
-      mujoco.mj_fullM(m, qm, d.qM)
+      mujoco.mju_sym2dense(qm, d.M, m.M_rownnz, m.M_rowadr, m.M_colind)
       # mjwarp adds padding to qM
       tu.assert_eq(qm, dx._impl.qM[: m.nv, : m.nv], 'qM')
       # qLD is fused in a cholesky factorize and solve, and not written to.