Remove dense code path for tendon Jacobian

PiperOrigin-RevId: 872444531
Change-Id: I6180101abc49469a72aee8bec5726e1e94f142ec

mjx/mujoco/mjx/_src/io.py

@@ -1111,7 +1111,7 @@ def _put_data_jax(
   impl_fields['actuator_moment'] = moment
 
   # convert ten_J to dense matrix
-  if mujoco.mj_isSparse(m):
+  if m.ntendon:
     ten_J = np.zeros((m.ntendon, m.nv))
     mujoco.mju_sparse2dense(
         ten_J,
@@ -1120,8 +1120,6 @@ def _put_data_jax(
         d.ten_J_rowadr,
         d.ten_J_colind,
     )
-  elif m.ntendon:
-    ten_J = d.ten_J.reshape((m.ntendon, m.nv))
   else:
     ten_J = np.zeros((m.ntendon, m.nv))
   impl_fields['ten_J'] = ten_J

mjx/mujoco/mjx/_src/smooth_test.py

@@ -121,8 +121,14 @@ class SmoothTest(absltest.TestCase):
     mujoco.mj_forward(m, d)
     # tendon
     dx = jax.jit(mjx.tendon)(mx, mjx.put_data(m, d))
-    # with dense jacobian mode, d.ten_J is already dense (ntendon*nv,), just reshape
-    ten_J = d.ten_J.reshape((m.ntendon, m.nv))
+    ten_J = np.zeros((m.ntendon, m.nv))
+    mujoco.mju_sparse2dense(
+        ten_J,
+        d.ten_J,
+        d.ten_J_rownnz,
+        d.ten_J_rowadr,
+        d.ten_J_colind,
+    )
     _assert_eq(ten_J, dx._impl.ten_J, 'ten_J')
     _assert_attr_eq(d, dx, 'ten_length')
     # transmission
@@ -397,19 +403,14 @@ class TendonTest(parameterized.TestCase):
     dx = jax.jit(mjx.forward)(mx, dx)
 
     _assert_eq(d.ten_length, dx.ten_length, 'ten_length')
-    # convert ten_J for comparison based on jacobian mode
-    if mujoco.mj_isSparse(m):
-      ten_J = np.zeros((m.ntendon, m.nv))
-      mujoco.mju_sparse2dense(
-          ten_J,
-          d.ten_J,
-          d.ten_J_rownnz,
-          d.ten_J_rowadr,
-          d.ten_J_colind,
-      )
-    else:
-      # dense mode: just reshape
-      ten_J = d.ten_J.reshape((m.ntendon, m.nv))
+    ten_J = np.zeros((m.ntendon, m.nv))
+    mujoco.mju_sparse2dense(
+        ten_J,
+        d.ten_J,
+        d.ten_J_rownnz,
+        d.ten_J_rowadr,
+        d.ten_J_colind,
+    )
     _assert_eq(ten_J, dx._impl.ten_J, 'ten_J')
     _assert_eq(d.ten_wrapnum, dx._impl.ten_wrapnum, 'ten_wrapnum')
     _assert_eq(d.ten_wrapadr, dx._impl.ten_wrapadr, 'ten_wrapadr')

mjx/mujoco/mjx/warp/forward_test.py

@@ -149,7 +149,14 @@ class ForwardTest(parameterized.TestCase):
         tu.assert_attr_eq(dx, d, 'cam_xpos')
         tu.assert_eq(dx.cam_xmat, d.cam_xmat.reshape((-1, 3, 3)), 'cam_xmat')
       tu.assert_attr_eq(dx, d, 'ten_length')
-      ten_J = d.ten_J.reshape((m.ntendon, m.nv))
+      ten_J = np.zeros((m.ntendon, m.nv))
+      mujoco.mju_sparse2dense(
+          ten_J,
+          d.ten_J,
+          d.ten_J_rownnz,
+          d.ten_J_rowadr,
+          d.ten_J_colind,
+      )
       tu.assert_eq(dx._impl.ten_J, ten_J, 'ten_J')
       tu.assert_attr_eq(dx._impl, d, 'ten_wrapadr')
       tu.assert_attr_eq(dx._impl, d, 'ten_wrapnum')