From 0a4eda50654be67f9bd7fcc4abaa32fa53b06cac Mon Sep 17 00:00:00 2001
From: Silvia Cruciani <silviacruciani@google.com>
Date: Mon, 17 Feb 2025 05:51:54 -0800
Subject: [PATCH] Fix setting slice of multidimensional arrays in mjx bind

PiperOrigin-RevId: 727834258
Change-Id: Id8fa25ab3aacb02603f2e0ec701f89bb69a7ac7e
---
 mjx/mujoco/mjx/_src/support.py      | 7 ++++---
 mjx/mujoco/mjx/_src/support_test.py | 4 ++++
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/mjx/mujoco/mjx/_src/support.py b/mjx/mujoco/mjx/_src/support.py
index 9c7c60cd..ea777a89 100644
--- a/mjx/mujoco/mjx/_src/support.py
+++ b/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})
 
diff --git a/mjx/mujoco/mjx/_src/support_test.py b/mjx/mujoco/mjx/_src/support_test.py
index 582d1c51..e4ae3d50 100644
--- a/mjx/mujoco/mjx/_src/support_test.py
+++ b/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):