From cf1b28d19a8428bed44633fee1b75c090f4292a7 Mon Sep 17 00:00:00 2001 From: Tarik Kelestemur Date: Wed, 27 May 2026 17:37:33 -0400 Subject: [PATCH] address btaba's review --- mjx/mujoco/mjx/_src/io_test.py | 28 ---------------------------- 1 file changed, 28 deletions(-) diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py index e086c38c..eba94658 100644 --- a/mjx/mujoco/mjx/_src/io_test.py +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -227,34 +227,6 @@ class ModelIOTest(parameterized.TestCase): m.light_active.astype(np.asarray(mx.light_active).dtype), ) - def test_replace_light_active(self): - """Tests that light_active can be replaced like other public fields.""" - m = mujoco.MjModel.from_xml_string(_LIGHTS) - mx = mjx.put_model(m, impl='jax') - - light_active = jp.array([True, False]) - mx_replaced = mx.replace(light_active=light_active) - - np.testing.assert_array_equal( - np.asarray(mx_replaced.light_active), np.asarray(light_active) - ) - np.testing.assert_array_equal(np.asarray(mx.light_active), m.light_active) - - def test_vmap_light_active_in_axes(self): - """Tests that light_active participates in batched model pytrees.""" - m = mujoco.MjModel.from_xml_string(_LIGHTS) - mx = mjx.put_model(m, impl='jax') - light_active = jp.array([[True, False], [False, True], [True, True]]) - mx_batched = mx.replace(light_active=light_active) - - in_axes = jax.tree_util.tree_map(lambda x: None, mx) - in_axes = in_axes.replace(light_active=0) - mapped = jax.vmap(lambda model: model.light_active, in_axes=(in_axes,))( - mx_batched - ) - - np.testing.assert_array_equal(np.asarray(mapped), np.asarray(light_active)) - def test_fluid_params(self): """Test that has_fluid_params is set when fluid params are present.""" m = mjx.put_model(