imt-ring 1.6.36__tar.gz → 1.6.38__tar.gz

{imt_ring-1.6.36 → imt_ring-1.6.38}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: imt-ring
-Version: 1.6.36
+Version: 1.6.38
 Summary: RING: Recurrent Inertial Graph-based Estimator
 Author-email: Simon Bachhuber <simon.bachhuber@fau.de>
 Project-URL: Homepage, https://github.com/SimiPixel/ring

{imt_ring-1.6.36 → imt_ring-1.6.38}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "imt-ring"
-version = "1.6.36"
+version = "1.6.38"
 authors = [
   { name="Simon Bachhuber", email="simon.bachhuber@fau.de" },
 ]

{imt_ring-1.6.36 → imt_ring-1.6.38}/src/imt_ring.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: imt-ring
-Version: 1.6.36
+Version: 1.6.38
 Summary: RING: Recurrent Inertial Graph-based Estimator
 Author-email: Simon Bachhuber <simon.bachhuber@fau.de>
 Project-URL: Homepage, https://github.com/SimiPixel/ring

{imt_ring-1.6.36 → imt_ring-1.6.38}/src/ring/algorithms/custom_joints/suntay.py

@@ -184,7 +184,7 @@ def register_suntay(sconfig: SuntayConfig, name: str = "suntay"):
 
         suntay_link_name = _utils_find_suntay_joint(sys)
 
-        params = jax.tree_map(
+        params = jax.tree.map(
             lambda arr: arr[sys.idx_map("l")[suntay_link_name]],
             sys.links.joint_params[name],
         )

{imt_ring-1.6.36 → imt_ring-1.6.38}/src/ring/algorithms/generator/batch.py

@@ -80,11 +80,11 @@ def generators_eager(
             # converts also to numpy; but with np.array.flags.writeable = False
             sample = jax.device_get(sample)
             # this then sets this flag to True
-            sample = jax.tree_map(np.array, sample)
+            sample = jax.tree.map(np.array, sample)
 
             sample_flat, _ = jax.tree_util.tree_flatten(sample)
             size = 1 if len(sample_flat) == 0 else sample_flat[0].shape[0]
-            callback([jax.tree_map(lambda a: a[i].copy(), sample) for i in range(size)])
+            callback([jax.tree.map(lambda a: a[i].copy(), sample) for i in range(size)])
 
             # cleanup
             del sample, sample_flat

{imt_ring-1.6.36 → imt_ring-1.6.38}/src/ring/algorithms/generator/finalize_fns.py

@@ -311,7 +311,7 @@ def _expand_then_flatten(Xy):
 
     X, y = _flatten(X), _flatten(y)
     if not batched:
-        X, y = jax.tree_map(lambda arr: arr[0], (X, y))
+        X, y = jax.tree.map(lambda arr: arr[0], (X, y))
     return X, y
 
 

{imt_ring-1.6.36 → imt_ring-1.6.38}/src/ring/algorithms/generator/pd_control.py

@@ -86,7 +86,7 @@ def _pd_control(P: jax.Array, D: Optional[jax.Array] = None):
         controller_state: PDControllerState, sys: base.System, state: base.State
     ) -> jax.Array:
         taus = jnp.zeros((sys.qd_size()))
-        q_ref, qd_ref = jax.tree_map(
+        q_ref, qd_ref = jax.tree.map(
             lambda arr: jax.lax.dynamic_index_in_dim(
                 arr, controller_state.i, keepdims=False
             ),

{imt_ring-1.6.36 → imt_ring-1.6.38}/src/ring/algorithms/kinematics.py

@@ -4,6 +4,7 @@ import jax
 import jax.numpy as jnp
 import jaxopt
 from jaxopt._src.base import Solver
+
 from ring import algebra
 from ring import base
 from ring import maths
@@ -171,7 +172,7 @@ def inverse_kinematics_endeffector(
 
             # find result of best q0 initial value
             best_q_index = jnp.argmin(values)
-            best_q, best_q_value = jax.tree_map(
+            best_q, best_q_value = jax.tree.map(
                 lambda arr: jax.lax.dynamic_index_in_dim(
                     arr, best_q_index, keepdims=False
                 ),

{imt_ring-1.6.36 → imt_ring-1.6.38}/src/ring/algorithms/sensors.py

@@ -244,7 +244,7 @@ def imu(
         measurements["mag"] = magnetometer(xs.rot, magvec)
 
     if smoothen_degree is not None:
-        measurements = jax.tree_map(
+        measurements = jax.tree.map(
             lambda arr: _moving_average(arr, smoothen_degree),
             measurements,
         )
@@ -257,7 +257,7 @@ def imu(
             delay = half_window
 
     if delay is not None and delay > 0:
-        measurements = jax.tree_map(
+        measurements = jax.tree.map(
             lambda arr: (jnp.pad(arr, ((delay, 0), (0, 0)))[:-delay]), measurements
         )
 
@@ -473,7 +473,7 @@ def _joint_axes_from_sys(sys: base.Transform, N: int) -> dict:
         X[name] = {"joint_axes": joint_axes}
 
     sys.scan(f, "lll", sys.link_names, sys.link_types, sys.links)
-    X = jax.tree_map(lambda arr: jnp.repeat(arr[None], N, axis=0), X)
+    X = jax.tree.map(lambda arr: jnp.repeat(arr[None], N, axis=0), X)
     return X
 
 
@@ -498,12 +498,12 @@ _quasi_physical_sys_str = r"""
 <x_xy>
     <options gravity="0 0 0"/>
     <worldbody>
-        <body name="IMU" joint="p3d" damping="0.1 0.1 0.1" spring_stiff="3 3 3">
-            <geom type="box" mass="0.002" dim="0.01 0.01 0.01"/>
+        <body name="IMU" joint="free" damping="1 1 1 10 10 10" spring_stiff="20 20 20 500 500 500">
+            <geom type="box" mass="1" dim="0.01 0.01 0.01"/>
         </body>
     </worldbody>
 </x_xy>
-"""
+"""  # noqa: E501
 
 
 def _quasi_physical_simulation_beautiful(
@@ -512,12 +512,14 @@ def _quasi_physical_simulation_beautiful(
     sys = io.load_sys_from_str(_quasi_physical_sys_str).replace(dt=dt)
 
     def step_dynamics(state: base.State, x):
-        state = algorithms.step(sys.replace(link_spring_zeropoint=x.pos), state)
+        state = algorithms.step(
+            sys.replace(link_spring_zeropoint=jnp.concatenate((x.rot, x.pos))), state
+        )
         return state, state.q
 
-    state = base.State.create(sys, q=xs.pos[0])
-    _, pos = jax.lax.scan(step_dynamics, state, xs)
-    return xs.replace(pos=pos)
+    state = base.State.create(sys, q=jnp.concatenate((xs.rot[0], xs.pos[0])))
+    _, qs = jax.lax.scan(step_dynamics, state, xs)
+    return xs.replace(rot=qs[:, :4], pos=qs[:, 4:])
 
 
 _constants = {

{imt_ring-1.6.36 → imt_ring-1.6.38}/src/ring/base.py

@@ -807,7 +807,7 @@ class System(_Base):
         if exists:
             return cls.from_xml(path, seed=seed)
         else:
-            return cls.from_str(path_or_str)
+            return cls.from_str(path_or_str, seed=seed)
 
     def coordinate_vector_to_q(
         self,

{imt_ring-1.6.36 → imt_ring-1.6.38}/src/ring/io/xml/from_xml.py

@@ -252,7 +252,7 @@ def load_sys_from_str(xml_str: str, seed: int = 1) -> base.System:
 
     # numpy -> jax
     # we load using numpy in order to have float64 precision
-    sys = jax.tree_map(jax.numpy.asarray, sys)
+    sys = jax.tree.map(jax.numpy.asarray, sys)
 
     sys = jcalc._init_joint_params(jax.random.PRNGKey(seed), sys)
 

{imt_ring-1.6.36 → imt_ring-1.6.38}/src/ring/ml/base.py

@@ -13,13 +13,13 @@ from ring.utils import pickle_save
 def _to_3d(tree):
     if tree is None:
         return None
-    return jax.tree_map(lambda arr: arr[None], tree)
+    return jax.tree.map(lambda arr: arr[None], tree)
 
 
 def _to_2d(tree, i: int = 0):
     if tree is None:
         return None
-    return jax.tree_map(lambda arr: arr[i], tree)
+    return jax.tree.map(lambda arr: arr[i], tree)
 
 
 class AbstractFilter(ABC):

{imt_ring-1.6.36 → imt_ring-1.6.38}/src/ring/ml/ml_utils.py

@@ -161,7 +161,7 @@ def _flatten_convert_filter_nested_dict(
     metrices: NestedDict, filter_nan_inf: bool = True
 ):
     metrices = _flatten_dict(metrices)
-    metrices = jax.tree_map(_to_float_if_not_string, metrices)
+    metrices = jax.tree.map(_to_float_if_not_string, metrices)
 
     if not filter_nan_inf:
         return metrices
@@ -216,7 +216,7 @@ def save_model_tf(jax_func, path: str, *input, validate: bool = True):
     from jax.experimental import jax2tf
     import tensorflow as tf
 
-    signature = jax.tree_map(
+    signature = jax.tree.map(
         lambda arr: tf.TensorSpec(list(arr.shape), tf.float32), input
     )
 
@@ -241,7 +241,7 @@ def save_model_tf(jax_func, path: str, *input, validate: bool = True):
     if validate:
         output_jax = jax_func(*input)
         output_tf = tf.saved_model.load(path)(*input)
-        jax.tree_map(
+        jax.tree.map(
             lambda a1, a2: np.allclose(a1, a2, atol=1e-5, rtol=1e-5),
             output_jax,
             output_tf,

{imt_ring-1.6.36 → imt_ring-1.6.38}/src/ring/ml/ringnet.py

@@ -248,7 +248,7 @@ class RING(ml_base.AbstractFilter):
         params, state = self.forward_lam_factory(lam=lam).init(key, X)
 
         if bs is not None:
-            state = jax.tree_map(lambda arr: jnp.repeat(arr[None], bs, axis=0), state)
+            state = jax.tree.map(lambda arr: jnp.repeat(arr[None], bs, axis=0), state)
 
         return params, state
 

{imt_ring-1.6.36 → imt_ring-1.6.38}/src/ring/ml/train.py

@@ -50,7 +50,7 @@ def _build_step_fn(
         # this vmap maps along batch-axis, not time-axis
         # time-axis is handled by `metric_fn`
         pipe = lambda q, qhat: jnp.mean(jax.vmap(metric_fn)(q, qhat))
-        error_tree = jax.tree_map(pipe, y, yhat)
+        error_tree = jax.tree.map(pipe, y, yhat)
         return jnp.mean(tree_utils.batch_concat(error_tree, 0)), state
 
     @partial(
@@ -274,7 +274,7 @@ def _build_eval_fn(
             ), f"The metric identitifier {metric_name} is not unique"
 
             pipe = lambda q, qhat: reduce_fn(jax.vmap(jax.vmap(metric_fn))(q, qhat))
-            values.update({metric_name: jax.tree_map(pipe, y, yhat)})
+            values.update({metric_name: jax.tree.map(pipe, y, yhat)})
 
         return values
 

{imt_ring-1.6.36 → imt_ring-1.6.38}/src/ring/rendering/mujoco_render.py

@@ -10,8 +10,8 @@ _skybox = """<texture name="skybox" type="skybox" builtin="gradient" rgb1=".4 .6
 _skybox_white = """<texture name="skybox" type="skybox" builtin="gradient" rgb1="1 1 1" rgb2="1 1 1" width="800" height="800" mark="random" markrgb="1 1 1"/>"""  # noqa: E501
 
 
-def _floor(floor_z: float) -> str:
-    return f"""<geom name="floor" pos="0 0 {floor_z}" size="0 0 1" type="plane" material="matplane" mass="0"/>"""  # noqa: E501
+def _floor(z: float, material: str) -> str:
+    return f"""<geom name="floor" pos="0 0 {z}" size="0 0 1" type="plane" material="{material}" mass="0"/>"""  # noqa: E501
 
 
 def _build_model_of_geoms(
@@ -19,7 +19,7 @@ def _build_model_of_geoms(
     cameras: dict[int, Sequence[str]],
     lights: dict[int, Sequence[str]],
     floor: bool,
-    floor_z: float,
+    floor_kwargs: dict,
     stars: bool,
     debug: bool,
 ) -> mujoco.MjModel:
@@ -77,10 +77,13 @@ def _build_model_of_geoms(
     xml_str = f""" # noqa: E501
 <mujoco>
   <asset>
-    <texture name="texplane" type="2d" builtin="checker" rgb1=".25 .25 .25" rgb2=".3 .3 .3" width="512" height="512" mark="cross" markrgb=".8 .8 .8"/>
+    <texture name="texplane" type="2d" builtin="checker" rgb1=".25 .25 .25" rgb2=".3 .3 .3" width="512" height="512" mark="cross" markrgb=".3 .3 .3"/>
     <material name="matplane" reflectance="0.3" texture="texplane" texrepeat="1 1" texuniform="true"/>
     <texture type="2d" name="groundplane" builtin="checker" mark="edge" rgb1="0.2 0.3 0.4" rgb2="0.1 0.2 0.3" markrgb="0.8 0.8 0.8" width="300" height="300"/>
     <material name="groundplane" texture="groundplane" texuniform="true" texrepeat="2 2" reflectance="0.2"/>
+    <material name="beige" rgba="0.76 0.80 0.50 1.0" specular="0.3" shininess="0.1" />
+    <material name="white" rgba="0.9 0.9 0.9 1.0" reflectance="0"/>
+    <material name="gray" rgba="0.4 0.5 0.5 1.0" reflectance="0.25"/>
     {_skybox if stars else ''}
     <texture name="grid" type="2d" builtin="checker" rgb1=".1 .2 .3" rgb2=".2 .3 .4" width="300" height="300" mark="edge" markrgb=".2 .3 .4"/>
     <material name="grid" texture="grid" texrepeat="1 1" texuniform="true" reflectance=".2"/>
@@ -98,7 +101,7 @@ def _build_model_of_geoms(
 <camera pos="0 -1 1" name="target" mode="targetbodycom" target="{targetbody}"/>
 <camera pos="0 -3 3" name="targetfar" mode="targetbodycom" target="{targetbody}"/>
 <camera pos="0 -5 5" name="targetFar" mode="targetbodycom" target="{targetbody}"/>
-{_floor(floor_z) if floor else ''}
+{_floor(**floor_kwargs) if floor else ''}
 {inside_worldbody_cameras}
 {inside_worldbody_lights}
 {inside_worldbody}
@@ -176,6 +179,7 @@ class MujocoScene:
         show_stars: bool = True,
         show_floor: bool = True,
         floor_z: float = -0.84,
+        floor_material: str = "matplane",
         debug: bool = False,
     ) -> None:
         self.debug = debug
@@ -190,7 +194,7 @@ class MujocoScene:
         self.add_cameras, self.add_lights = to_list(add_cameras), to_list(add_lights)
         self.show_stars = show_stars
         self.show_floor = show_floor
-        self.floor_z = floor_z
+        self.floor_kwargs = dict(z=floor_z, material=floor_material)
 
     def init(self, geoms: list[base.Geometry]):
         self._parent_ids = list(set([geom.link_idx for geom in geoms]))
@@ -199,7 +203,7 @@ class MujocoScene:
             self.add_cameras,
             self.add_lights,
             floor=self.show_floor,
-            floor_z=self.floor_z,
+            floor_kwargs=self.floor_kwargs,
             stars=self.show_stars,
             debug=self.debug,
         )

{imt_ring-1.6.36 → imt_ring-1.6.38}/src/ring/rendering/vispy_render.py

@@ -7,14 +7,15 @@ from typing import Optional, TypeVar
 import jax
 import jax.numpy as jnp
 import numpy as np
-from ring import algebra
-from ring import base
-from ring import maths
 from tree_utils import PyTree
 from tree_utils import tree_batch
 from vispy import scene
 from vispy.scene import MatrixTransform
 
+from ring import algebra
+from ring import base
+from ring import maths
+
 from . import vispy_visuals
 
 Camera = TypeVar("Camera")
@@ -192,7 +193,7 @@ class Scene(ABC):
 
         # step 3: update visuals
         for i, (visual, geom) in enumerate(zip(self.visuals, self.geoms)):
-            t = jax.tree_map(lambda arr: arr[i], transform_per_visual)
+            t = jax.tree.map(lambda arr: arr[i], transform_per_visual)
             if self._fresh_init:
                 self._init_visual(visual, t, geom)
             else:

{imt_ring-1.6.36 → imt_ring-1.6.38}/src/ring/sys_composer/inject_sys.py

@@ -2,12 +2,13 @@ from typing import Optional
 
 import jax
 import jax.numpy as jnp
-from ring import base
 from tree_utils import tree_batch
 
+from ring import base
+
 
 def _tree_nan_like(tree, repeats: int):
-    return jax.tree_map(
+    return jax.tree.map(
         lambda arr: jnp.repeat(arr[0:1] * jnp.nan, repeats, axis=0), tree
     )
 

{imt_ring-1.6.36 → imt_ring-1.6.38}/src/ring/utils/batchsize.py

@@ -39,19 +39,19 @@ def merge_batchsize(
     tree: PyTree, pmap_size: int, vmap_size: int, third_dim_also: bool = False
 ) -> PyTree:
     if third_dim_also:
-        return jax.tree_map(
+        return jax.tree.map(
             lambda arr: arr.reshape(
                 (pmap_size * vmap_size * arr.shape[2],) + arr.shape[3:]
             ),
             tree,
         )
-    return jax.tree_map(
+    return jax.tree.map(
         lambda arr: arr.reshape((pmap_size * vmap_size,) + arr.shape[2:]), tree
     )
 
 
 def expand_batchsize(tree: PyTree, pmap_size: int, vmap_size: int) -> PyTree:
-    return jax.tree_map(
+    return jax.tree.map(
         lambda arr: arr.reshape(
             (
                 pmap_size,

{imt_ring-1.6.36 → imt_ring-1.6.38}/src/ring/utils/dataloader.py

@@ -4,14 +4,15 @@ from typing import Callable, Optional
 
 import jax
 import numpy as np
-from ring.utils import parse_path
-from ring.utils import pickle_load
 import torch
 from torch.utils.data import DataLoader
 from torch.utils.data import Dataset
 import tqdm
 from tree_utils import PyTree
 
+from ring.utils import parse_path
+from ring.utils import pickle_load
+
 
 def make_generator(
     *paths,
@@ -103,7 +104,7 @@ def pytorch_generator(
     dl_iter = iter(dl)
 
     def to_numpy(tree: PyTree[torch.Tensor]):
-        return jax.tree_map(lambda tensor: tensor.numpy(), tree)
+        return jax.tree.map(lambda tensor: tensor.numpy(), tree)
 
     def generator(_):
         nonlocal dl, dl_iter

{imt_ring-1.6.36 → imt_ring-1.6.38}/src/ring/utils/dataloader_torch.py

@@ -1,16 +1,25 @@
 import os
+import pickle
 from typing import Any, Optional
 import warnings
 
-import jax
 import numpy as np
 import torch
 from torch.utils.data import DataLoader
 from torch.utils.data import Dataset
+import tree
 from tree_utils import PyTree
 
-from ring.utils import parse_path
-from ring.utils import pickle_load
+from ring.utils.path import parse_path
+
+
+def pickle_load(
+    path,
+):
+    path = parse_path(path, extension="pickle", require_is_file=True)
+    with open(path, "rb") as file:
+        obj = pickle.load(file)
+    return obj
 
 
 class FolderOfFilesDataset(Dataset):
@@ -60,8 +69,8 @@ def dataset_to_generator(
     )
     dl_iter = iter(dl)
 
-    def to_numpy(tree: PyTree[torch.Tensor]):
-        return jax.tree_map(lambda tensor: tensor.numpy(), tree)
+    def to_numpy(data: PyTree[torch.Tensor]):
+        return tree.map_structure(lambda tensor: tensor.numpy(), data)
 
     def generator(_):
         nonlocal dl, dl_iter

{imt_ring-1.6.36 → imt_ring-1.6.38}/src/ring/utils/hdf5.py

@@ -121,7 +121,7 @@ def _parse_path(
 
 def _tree_concat(trees: list):
     # otherwise scalar-arrays will lead to indexing error
-    trees = jax.tree_map(lambda arr: np.atleast_1d(arr), trees)
+    trees = jax.tree.map(lambda arr: np.atleast_1d(arr), trees)
 
     if len(trees) == 0:
         return trees

{imt_ring-1.6.36 → imt_ring-1.6.38}/src/ring/utils/normalizer.py

@@ -3,9 +3,10 @@ from typing import Callable, TypeVar
 
 import jax
 import jax.numpy as jnp
-from ring.algorithms.generator import types
 import tree_utils
 
+from ring.algorithms.generator import types
+
 KEY = jax.random.PRNGKey(777)
 KEY_PERMUTATION = jax.random.PRNGKey(888)
 
@@ -37,12 +38,12 @@ def make_normalizer_from_generator(
     # permute 0-th axis, since batchsize of generator might be larger than
     # `approx_with_large_batchsize`, then we would not get a representative
     # subsample otherwise
-    Xs = jax.tree_map(lambda arr: jax.random.permutation(KEY_PERMUTATION, arr), Xs)
+    Xs = jax.tree.map(lambda arr: jax.random.permutation(KEY_PERMUTATION, arr), Xs)
     Xs = tree_utils.tree_slice(Xs, start=0, slice_size=approx_with_large_batchsize)
 
     # obtain statistics
-    mean = jax.tree_map(lambda arr: jnp.mean(arr, axis=(0, 1)), Xs)
-    std = jax.tree_map(lambda arr: jnp.std(arr, axis=(0, 1)), Xs)
+    mean = jax.tree.map(lambda arr: jnp.mean(arr, axis=(0, 1)), Xs)
+    std = jax.tree.map(lambda arr: jnp.std(arr, axis=(0, 1)), Xs)
 
     if verbose:
         print("Mean: ", mean)
@@ -51,6 +52,6 @@ def make_normalizer_from_generator(
     eps = 1e-8
 
     def normalizer(X):
-        return jax.tree_map(lambda a, b, c: (a - b) / (c + eps), X, mean, std)
+        return jax.tree.map(lambda a, b, c: (a - b) / (c + eps), X, mean, std)
 
     return normalizer

{imt_ring-1.6.36 → imt_ring-1.6.38}/src/ring/utils/utils.py

@@ -1,5 +1,6 @@
 from importlib import import_module as _import_module
 import io
+from pathlib import Path
 import pickle
 import random
 from typing import Optional
@@ -152,13 +153,28 @@ def import_lib(
 
 def pickle_save(obj, path, overwrite: bool = False):
     path = parse_path(path, extension="pickle", file_exists_ok=overwrite)
-    with open(path, "wb") as file:
-        pickle.dump(obj, file, protocol=5)
+    try:
+        with open(path, "wb") as file:
+            pickle.dump(obj, file, protocol=5)
+    except OSError as e:
+        print(
+            f"saving with `pickle` throws exception {e}. "
+            + "Attempting to save using `joblib`"
+        )
+        path = parse_path(path, extension="joblib", file_exists_ok=overwrite)
+        import joblib
+
+        joblib.dump(obj, path)
 
 
 def pickle_load(
     path,
 ):
+    if Path(path).suffix == ".joblib":
+        import joblib
+
+        return joblib.load(path)
+
     path = parse_path(path, extension="pickle", require_is_file=True)
     with open(path, "rb") as file:
         obj = pickle.load(file)

{imt_ring-1.6.36 → imt_ring-1.6.38}/tests/test_ml_utils.py

@@ -41,7 +41,7 @@ def test_save_load_generators():
     data = rcmg.to_list()[0]
     rcmg.to_pickle(path)
 
-    data_list = [jax.tree_map(lambda a: a[0], utils.pickle_load(path))]
+    data_list = [jax.tree.map(lambda a: a[0], utils.pickle_load(path))]
     gen_reloaded = ring.RCMG.eager_gen_from_list(data_list, 1)
     data_reloaded = unbatch_gen(gen_reloaded)(jax.random.PRNGKey(1))
 

{imt_ring-1.6.36 → imt_ring-1.6.38}/tests/test_sim2real.py

@@ -2,6 +2,7 @@ from _compat import unbatch_gen
 import jax
 import jax.numpy as jnp
 import numpy as np
+
 import ring
 from ring import maths
 from ring import sim2real
@@ -49,7 +50,7 @@ def test_forward_kinematics_omc():
             # t1_omc should be used when p == -1, else t1_sys
             @jax.vmap
             def merge_transform1(t1_omc):
-                return jax.tree_map(
+                return jax.tree.map(
                     lambda a, b: jnp.where(
                         jnp.repeat(
                             jnp.array(sys.link_parents)[:, None] == -1,
@@ -138,7 +139,7 @@ def test_zip_unzip_scale():
         t1, t2 = sim2real.unzip_xs(sys, xs)
         xs_re = sim2real.zip_xs(sys, t1, t2)
 
-        jax.tree_map(
+        jax.tree.map(
             lambda a, b: np.testing.assert_allclose(a, b, rtol=1e-3, atol=1e-5),
             xs,
             xs_re,