chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Small fix to testspeed

Browse Source 
PiperOrigin-RevId: 840607453
Change-Id: If913866cc0a3d4d3e6c39c7b5948fa167678ccad
This commit is contained in:
Baruch Tabanpour
2025-12-05 01:27:36 -08:00
committed by Copybara-Service
parent 237a5e5c55
commit fc109f0ab5
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/warp/testspeed.py
+3 -3
View File
@@ -93,7 +93,7 @@ def benchmark(
@jax.vmap @jax.vmap
def init(key): def init(key):
d = mjx.make_data( 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 return d
@@ -262,7 +262,7 @@ def benchmark_raw_jax_warp(
d_ = mujoco.MjData(m) d_ = mujoco.MjData(m)
mw = mjwarp.put_model(m) mw = mjwarp.put_model(m)
d = mjwarp.put_data( 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) jax_unroll_fn = jax_jit(unroll)
@@ -310,7 +310,7 @@ def benchmark_raw_warp(
mw = mjwarp.put_model(m) mw = mjwarp.put_model(m)
dw = mjwarp.make_data( 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': if function == 'kinematics':