Move flexedge sparse Jacobian indices to mjModel.

PiperOrigin-RevId: 855718305
Change-Id: Ic136f4efa6735b934288f65c332d71ee8b0f6da5

mjx/mujoco/mjx/_src/io.py

@@ -754,9 +754,6 @@ def _make_data_c(
       'light_xdir': (m.nlight, 3, float_),
       'flexvert_xpos': (nflexvert, 3, float_),
       'flexelem_aabb': (nflexelem, 6, float_),
-      'flexedge_J_rownnz': (nflexedge, np.int32),
-      'flexedge_J_rowadr': (nflexedge, np.int32),
-      'flexedge_J_colind': (nflexedge, m.nv, np.int32),
       'flexedge_J': (nflexedge, m.nv, float_),
       'flexedge_length': (nflexedge, float_),
       'ten_J_rownnz': (m.ntendon, np.int32),

mjx/mujoco/mjx/_src/types.py

@@ -626,6 +626,9 @@ class ModelC(PyTreeNode):
   flex_centered: jax.Array
   flex_bvhadr: jax.Array
   flex_bvhnum: jax.Array
+  flexedge_J_rownnz: jax.Array
+  flexedge_J_rowadr: jax.Array
+  flexedge_J_colind: jax.Array
   mesh_polynum: jax.Array
   mesh_polyadr: jax.Array
   mesh_polynormal: jax.Array
@@ -1047,9 +1050,6 @@ class DataC(PyTreeNode):
   cinert: jax.Array
   flexvert_xpos: jax.Array
   flexelem_aabb: jax.Array
-  flexedge_J_rownnz: jax.Array  # pylint:disable=invalid-name
-  flexedge_J_rowadr: jax.Array  # pylint:disable=invalid-name
-  flexedge_J_colind: jax.Array  # pylint:disable=invalid-name
   flexedge_J: jax.Array  # pylint:disable=invalid-name
   flexedge_length: jax.Array
   bvh_aabb_dyn: jax.Array

mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py

@@ -858,7 +858,7 @@ def put_data(
     d.ten_J = wp.array(np.full((nworld, mjm.ntendon, mjm.nv), ten_J), dtype=float)
     flexedge_J = np.zeros((mjm.nflexedge, mjm.nv))
     mujoco.mju_sparse2dense(
-      flexedge_J, mjd.flexedge_J.reshape(-1), mjd.flexedge_J_rownnz, mjd.flexedge_J_rowadr, mjd.flexedge_J_colind.reshape(-1)
+      flexedge_J, mjd.flexedge_J.reshape(-1), mjm.flexedge_J_rownnz, mjm.flexedge_J_rowadr, mjm.flexedge_J_colind.reshape(-1)
     )
     d.flexedge_J = wp.array(np.full((nworld, mjm.nflexedge, mjm.nv), flexedge_J), dtype=float)
   else: