Fix setting slice of multidimensional arrays in mjx bind

PiperOrigin-RevId: 727834258
Change-Id: Id8fa25ab3aacb02603f2e0ec701f89bb69a7ac7e

mjx/mujoco/mjx/_src/support.py

@@ -490,14 +490,15 @@ class BindData(object):
       typ = self.model.jnt_type[self.id]
       num = sum((typ == jt) * jt.dof_width() for jt in JointType)
     elif isinstance(self.id, list):
-      adr = self.id
+      adr = self.id * dim
       num = [dim for _ in range(len(self.id))]
     else:
-      adr = [self.id]
+      adr = [self.id * dim]
       num = [dim]
     i = 0
     for a, n in zip(adr, num):
-      array = array.at[a: a + n].set(value[i: i + n])
+      shape = array.shape
+      array = array.flatten().at[a : a + n].set(value[i : i + n]).reshape(shape)
       i += n
     return self.data.replace(**{self.__getname(name): array})
 

mjx/mujoco/mjx/_src/support_test.py

@@ -276,6 +276,10 @@ class SupportTest(parameterized.TestCase):
     np.testing.assert_array_equal(
         dx7.bind(mx, s.bodies[1]).xfrc_applied, [1, 2, 3, 4, 5, 6]
     )
+    for body in s.bodies[:1] + s.bodies[2:]:
+      np.testing.assert_array_equal(
+          dx7.bind(mx, body).xfrc_applied, [0, 0, 0, 0, 0, 0]
+      )
 
     # test invalid name
     with self.assertRaises(AttributeError):