Use jax.tree_util.tree_map in place of deprecated tree_multimap.

The latter is a simple alias of the former, so this change is a no-op.

PiperOrigin-RevId: 461229165

physics_inspired_models/integrators.py

@@ -204,7 +204,7 @@ def solve_ivp_dt(
     for t_and_dt_i in zip(t, dt):
       y.append(loop_body(y[-1], t_and_dt_i)[0])
     # Note that we do not return the initial point
-    return t_eval, jax.tree_multimap(lambda *args: jnp.stack(args, axis=0),
+    return t_eval, jax.tree_map(lambda *args: jnp.stack(args, axis=0),
                                      *y[1:])
 
 
@@ -252,7 +252,7 @@ def solve_ivp_dt_two_directions(
     )[1]
     yt.append(yt_fwd)
   if len(yt) > 1:
-    return jax.tree_multimap(lambda *a: jnp.concatenate(a, axis=0), *yt)
+    return jax.tree_map(lambda *a: jnp.concatenate(a, axis=0), *yt)
   else:
     return yt[0]
 

physics_inspired_models/jaxline_train.py

@@ -192,7 +192,7 @@ class HGNExperiment(experiment.AbstractExperiment):
     new_state = utils.pmean_if_pmap(new_state, axis_name="i")
     new_state = hk.data_structures.to_mutable_dict(new_state)
     new_state = hk.data_structures.to_immutable_dict(new_state)
-    return jax.tree_multimap(jnp.add, new_state, state)
+    return jax.tree_map(jnp.add, new_state, state)
 
   #                  _
   #   _____   ____ _| |

physics_inspired_models/models/dynamics.py

@@ -829,7 +829,7 @@ class DiscreteDynamicsNetwork(hk.Module):
       if len(yt) == 1:
         yt = yt[0][:, None]
       else:
-        yt = jax.tree_multimap(lambda args: jnp.stack(args, 1), yt)
+        yt = jax.tree_map(lambda args: jnp.stack(args, 1), yt)
     if return_stats:
       return yt, dict()
     else:

physics_inspired_models/utils.py

@@ -220,11 +220,11 @@ class MultiBatchAccumulator(object):
       self._obj = jax.tree_map(lambda y: y * num_samples, averaged_values)
       self._num_samples = num_samples
     else:
-      self._obj_max = jax.tree_multimap(jnp.maximum, self._obj_max,
+      self._obj_max = jax.tree_map(jnp.maximum, self._obj_max,
                                         averaged_values)
-      self._obj_min = jax.tree_multimap(jnp.minimum, self._obj_min,
+      self._obj_min = jax.tree_map(jnp.minimum, self._obj_min,
                                         averaged_values)
-      self._obj = jax.tree_multimap(lambda x, y: x + y * num_samples, self._obj,
+      self._obj = jax.tree_map(lambda x, y: x + y * num_samples, self._obj,
                                     averaged_values)
       self._num_samples += num_samples
 
@@ -249,7 +249,7 @@ register_pytree_node(
 
 
 def inner_product(x: Any, y: Any) -> jnp.ndarray:
-  products = jax.tree_multimap(lambda x_, y_: jnp.sum(x_ * y_), x, y)
+  products = jax.tree_map(lambda x_, y_: jnp.sum(x_ * y_), x, y)
   return sum(jax.tree_leaves(products))