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

Fix jax deprecation warning for jax.tree_map.

Browse Source 
PiperOrigin-RevId: 626092308
Change-Id: I78145efe62ae118726c1755856d7cb9b1abd988a
This commit is contained in:
Google DeepMind
2024-04-18 11:12:41 -07:00
committed by Copybara-Service
parent 4b7b40d329
commit 8b7f1094f1
15 changed files with 52 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/collision_sdf.py
+2 -2
View File
@@ -131,7 +131,7 @@ def _optim(
def capsule_ellipsoid(c: GeomInfo, e: GeomInfo) -> Contact:
""""Calculates contact between a capsule and an ellipsoid."""
pos, dist, n = _optim(_capsule, _ellipsoid, c, e)
return jax.tree_map(
return jax.tree_util.tree_map(
lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n))
)
@@ -139,7 +139,7 @@ def capsule_ellipsoid(c: GeomInfo, e: GeomInfo) -> Contact:
def ellipsoid_ellipsoid(e1: GeomInfo, e2: GeomInfo) -> Contact:
""""Calculates contact between two ellipsoids."""
pos, dist, n = _optim(_ellipsoid, _ellipsoid, e1, e2)
return jax.tree_map(
return jax.tree_util.tree_map(
lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n))
)