From fc109f0ab593d7a5bdb876479a1f122b9c823c30 Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Fri, 5 Dec 2025 01:27:36 -0800 Subject: [PATCH] Small fix to testspeed PiperOrigin-RevId: 840607453 Change-Id: If913866cc0a3d4d3e6c39c7b5948fa167678ccad --- mjx/mujoco/mjx/warp/testspeed.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mjx/mujoco/mjx/warp/testspeed.py b/mjx/mujoco/mjx/warp/testspeed.py index 8faba97f..fef2a386 100644 --- a/mjx/mujoco/mjx/warp/testspeed.py +++ b/mjx/mujoco/mjx/warp/testspeed.py @@ -93,7 +93,7 @@ def benchmark( @jax.vmap def init(key): d = mjx.make_data( - m, impl=mx.impl, nconmax=_NCONMAX.value, njmax=_NJMAX.value + m, impl=mx.impl, naconmax=_NACONMAX.value, njmax=_NJMAX.value ) return d @@ -262,7 +262,7 @@ def benchmark_raw_jax_warp( d_ = mujoco.MjData(m) mw = mjwarp.put_model(m) d = mjwarp.put_data( - m, d_, nworld=nenv, nconmax=_NCONMAX.value, njmax=_NJMAX.value + m, d_, nworld=nenv, naconmax=_NACONMAX.value, njmax=_NJMAX.value ) jax_unroll_fn = jax_jit(unroll) @@ -310,7 +310,7 @@ def benchmark_raw_warp( mw = mjwarp.put_model(m) dw = mjwarp.make_data( - m, nworld=nenv, nconmax=_NCONMAX.value, njmax=_NJMAX.value + m, nworld=nenv, nconmax=_NACONMAX.value, njmax=_NJMAX.value ) if function == 'kinematics':