Fix setting slices of multidimensional arrays for multiple targets

PiperOrigin-RevId: 758623598
Change-Id: Ie05b36c25be74ddf7451a95377caf1ed1d692141

mjx/mujoco/mjx/_src/support.py

@@ -499,12 +499,13 @@ class BindData(object):
         adr = [adr]
         num = [num]
     elif isinstance(self.id, list):
-      adr = self.id * dim
+      adr = (np.array(self.id) * dim).tolist()
       num = [dim for _ in range(len(self.id))]
     else:
       adr = [self.id * dim]
       num = [dim]
     i = 0
+    value = jax.numpy.array(value).flatten()
     for a, n in zip(adr, num):
       shape = array.shape
       array = array.flatten().at[a : a + n].set(value[i : i + n]).reshape(shape)

mjx/mujoco/mjx/_src/support_test.py

@@ -308,6 +308,16 @@ class SupportTest(parameterized.TestCase):
       np.testing.assert_array_equal(
           dx7.bind(mx, body).xfrc_applied, [0, 0, 0, 0, 0, 0]
       )
+    dx10 = dx.bind(mx, s.bodies[0:2]).set(
+        'xfrc_applied',
+        np.array([np.array([1, 1, 1, 1, 1, 1]), np.array([2, 2, 2, 2, 2, 2])]),
+    )
+    np.testing.assert_array_equal(
+        dx10.bind(mx, s.bodies[0]).xfrc_applied, [1, 1, 1, 1, 1, 1]
+    )
+    np.testing.assert_array_equal(
+        dx10.bind(mx, s.bodies[1]).xfrc_applied, [2, 2, 2, 2, 2, 2]
+    )
 
     # test attribute and type mismatches
     with self.assertRaisesRegex(