Remove disable_jit in test.

PiperOrigin-RevId: 614824680
Change-Id: I15f1441dfb1a70b98e5a221c0e15b06572860748
This commit is contained in:
Baruch Tabanpour
2024-03-11 16:02:31 -07:00
committed by Copybara-Service
parent 8413a615d7
commit a49f3f63fc
@@ -297,7 +297,6 @@ class ConvexTest(absltest.TestCase):
def test_flat_box_plane(self): def test_flat_box_plane(self):
"""Tests box collision with a plane.""" """Tests box collision with a plane."""
with jax.disable_jit():
d, dx = _collide(self._FLAT_BOX_PLANE) d, dx = _collide(self._FLAT_BOX_PLANE)
np.testing.assert_array_less(dx.contact.dist, 0) np.testing.assert_array_less(dx.contact.dist, 0)
@@ -586,7 +585,6 @@ class TopKContactTest(absltest.TestCase):
def test_max_pair(self): def test_max_pair(self):
"""Tests contact culling before the collision functions were dispatched.""" """Tests contact culling before the collision functions were dispatched."""
with jax.disable_jit():
m = mujoco.MjModel.from_xml_string(self._CAPSULES_MAX_PAIR) m = mujoco.MjModel.from_xml_string(self._CAPSULES_MAX_PAIR)
mx_top_k = mjx.put_model(m) mx_top_k = mjx.put_model(m)
mx_all = mx_top_k.replace( mx_all = mx_top_k.replace(