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, )