From c1d97d4a22512cde88deb81e2e8c7b194b64bfc1 Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Fri, 6 Dec 2024 10:45:37 -0800 Subject: [PATCH] Fix #2237. PiperOrigin-RevId: 703547316 Change-Id: I279876a299e56a16f8d5227fc1d0f64265477853 --- mjx/mujoco/mjx/_src/io.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 0988e25c..dad0b573 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -233,9 +233,11 @@ def make_data( solreffriction=jp.zeros((ncon, mujoco.mjNREF), dtype=float), solimp=jp.zeros((ncon, mujoco.mjNIMP), dtype=float), dim=dim, - geom1=jp.full((ncon,), -1, dtype=jp.int32), - geom2=jp.full((ncon,), -1, dtype=jp.int32), - geom=jp.full((ncon, 2), -1, dtype=jp.int32), + # let jax pick contact.geom int precision, for interop with + # jax_enable_x64 + geom1=jp.full((ncon,), -1, dtype=int), + geom2=jp.full((ncon,), -1, dtype=int), + geom=jp.full((ncon, 2), -1, dtype=int), efc_address=efc_address, )