Changing deprecated tree_multimap to tree_map.

PiperOrigin-RevId: 439855058

Author: botev

examples/optimizers.py

@@ -164,16 +164,16 @@ def tf1_rmsprop(
   def tf1_scale_by_rms(decay_=0.9, epsilon_=1e-8):
     """Same as optax.scale_by_rms, but initializes second moment to one."""
     def init_fn(params):
-      nu = jax.tree_multimap(jnp.ones_like, params)  # second moment
+      nu = jax.tree_map(jnp.ones_like, params)  # second moment
       return optax.ScaleByRmsState(nu=nu)
     def _update_moment(updates, moments, decay, order):
-      return jax.tree_multimap(
+      return jax.tree_map(
           lambda g, t: (1 - decay) * (g ** order) + decay * t, updates, moments)
     def update_fn(updates, state, params=None):
       del params
       nu = _update_moment(updates, state.nu, decay_, 2)
-      updates = jax.tree_multimap(lambda g, n: g / (jnp.sqrt(n + epsilon_)),
-                                  updates, nu)
+      updates = jax.tree_map(lambda g, n: g / (jnp.sqrt(n + epsilon_)),
+                             updates, nu)
       return updates, optax.ScaleByRmsState(nu=nu)
     return optax.GradientTransformation(init_fn, update_fn)
 

kfac_jax/_src/optimizer.py

@@ -805,7 +805,7 @@ def _step(
       update_norm = utils.norm(delta)
 
     # Update parameters
-    params = jax.tree_multimap(jnp.add, params, delta)
+    params = jax.tree_map(jnp.add, params, delta)
 
     # Optionally compute the reduction ratio and update the damping
     if self._use_adaptive_damping:

kfac_jax/_src/tag_graph_matcher.py

@@ -678,7 +678,7 @@ def write_env(
   if isinstance(var, list):
     if not isinstance(val, list):
       val = [val]
-    return jax.tree_multimap(lambda x, y: write_env(env, x, y), var, val)
+    return jax.tree_map(lambda x, y: write_env(env, x, y), var, val)
   elif isinstance(var, (core.Literal, core.Var)):
     env[var] = val
   else:

kfac_jax/_src/utils.py

@@ -160,7 +160,7 @@ def average_func(*args) -> PyTree:
       def scan_fn(accumulator, args_):
         vmap_value = vmap_fn(*args_)
         avg_value = jax.tree_map(lambda x: jnp.mean(x, axis=0), vmap_value)
-        return jax.tree_multimap(jnp.add, accumulator, avg_value), None
+        return jax.tree_map(jnp.add, accumulator, avg_value), None
 
       loop_shape = (num_parallel_chunks, parallel_size)
       loop_args = jax.tree_map(
@@ -435,7 +435,7 @@ def weighted_sum_of_objects(
     if not abstract_objects_equal(accumulator, o_i):
       raise ValueError("One or more objects do not have equivalent abstract "
                        "structure.")
-    accumulator = jax.tree_multimap(jnp.add, accumulator, scalar_mul(o_i, c_i))
+    accumulator = jax.tree_map(jnp.add, accumulator, scalar_mul(o_i, c_i))
   return accumulator
 
 
@@ -448,7 +448,7 @@ def array_ip(x, y):
     return jnp.dot(x, y, precision=lax.Precision.HIGHEST)
 
   with jax.experimental.enable_x64():
-    elements_inner_products = jax.tree_multimap(array_ip, obj1, obj2)
+    elements_inner_products = jax.tree_map(array_ip, obj1, obj2)
     flat_list = jax.tree_leaves(elements_inner_products)
     result = flat_list[0]
     for element_ip in flat_list[1:]:
@@ -484,7 +484,7 @@ def inner_product(
     raise ValueError("The objects do not have identical abstract structure.")
   if in_float64:
     return _inner_product_float64(obj1, obj2)
-  elements_product = jax.tree_multimap(lambda x, y: jnp.sum(x * y), obj1, obj2)
+  elements_product = jax.tree_map(lambda x, y: jnp.sum(x * y), obj1, obj2)
   return sum(jax.tree_leaves(elements_product))
 
 
@@ -587,7 +587,7 @@ def block_permuted(
 
 def norm(obj: PyTree) -> chex.Array:
   """Computes the Euclidean norm of the provided PyTree object."""
-  elements_squared_norm = jax.tree_multimap(
+  elements_squared_norm = jax.tree_map(
       lambda x: jnp.sum(jnp.square(x)), obj)
   return jnp.sqrt(sum(jax.tree_flatten(elements_squared_norm)[0]))
 
@@ -985,7 +985,7 @@ def add(self, value_obj: PyTree, weight: chex.Numeric = 1) -> None:
         raise ValueError("The provided `value_obj` has an empty PyTree "
                          "structure, but the accumulator has been initialized "
                          "with a non-empty PyTree object.")
-      self._accumulator = jax.tree_multimap(
+      self._accumulator = jax.tree_map(
           jnp.add, self._accumulator, value_obj)
     elif not tree_is_empty(value_obj):
       raise ValueError("The provided `value_obj` has a non-empty PyTree "
@@ -1261,7 +1261,7 @@ def decorated(instance: "WithStagedMethods", *args: Any) -> PyTree:
           ]
           outs.append(method(instance, *args_i))
 
-        outs = jax.tree_multimap(jnp.stack, *outs)
+        outs = jax.tree_map(jnp.stack, *outs)
 
       elif instance.debug:
         outs = method(instance, *args)

tests/test_tracer.py

@@ -102,10 +102,10 @@ def compare_multi_batch(
     for d in (data1, data2):
       outputs.append(func(d))
     if combine == "concatenate":
-      outputs = jax.tree_multimap(
+      outputs = jax.tree_map(
           lambda x, y: jnp.concatenate([x, y], axis=0), *outputs)
     elif combine == "sum":
-      outputs = jax.tree_multimap(lambda x, y: x + y, *outputs)
+      outputs = jax.tree_map(lambda x, y: x + y, *outputs)
     else:
       raise NotImplementedError()