Fix jax deprecation warning for jax.tree_map.

PiperOrigin-RevId: 626092308 Change-Id: I78145efe62ae118726c1755856d7cb9b1abd988a
2024-04-18 11:12:41 -07:00
parent 4b7b40d329
commit 8b7f1094f1
15 changed files with 52 additions and 52 deletions
@@ -829,7 +829,7 @@
    "\n",
    "  friction, gain, bias = rand(rng)\n",
    "\n",
-    "  in_axes = jax.tree_map(lambda x: None, sys)\n",
+    "  in_axes = jax.tree_util.tree_map(lambda x: None, sys)\n",
    "  in_axes = in_axes.tree_replace({\n",
    "      'geom_friction': 0,\n",
    "      'actuator_gainprm': 0,\n",