From 4ba04586ae78fe811fac14e3ba48c14469aeda85 Mon Sep 17 00:00:00 2001 From: Silvia Cruciani Date: Wed, 14 May 2025 04:35:56 -0700 Subject: [PATCH] Fix setting slices of multidimensional arrays for multiple targets PiperOrigin-RevId: 758623598 Change-Id: Ie05b36c25be74ddf7451a95377caf1ed1d692141 --- mjx/mujoco/mjx/_src/support.py | 3 ++- mjx/mujoco/mjx/_src/support_test.py | 10 ++++++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/mjx/mujoco/mjx/_src/support.py b/mjx/mujoco/mjx/_src/support.py index d68c0407..aea5bced 100644 --- a/mjx/mujoco/mjx/_src/support.py +++ b/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) diff --git a/mjx/mujoco/mjx/_src/support_test.py b/mjx/mujoco/mjx/_src/support_test.py index 5f060bc2..38392d38 100644 --- a/mjx/mujoco/mjx/_src/support_test.py +++ b/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(