Merge pull request #2026 from jheek:remove-more-tree_multimap

PiperOrigin-RevId: 439524627

Author: Flax Authors

docs/flip/1009-optimizer-api.md

@@ -349,7 +349,7 @@ Remarks:
   `OptimizerDef`).
 - The functions `init_param_state()` and `apply_param_gradient()` are called
   for every leaf in the params/grads pytree. This makes it possible to write the
-  calculations directly without `jax.tree_multimap()`.
+  calculations directly without `jax.tree_map()`.
 - The interface was defined in pre-Linen without the distinction of `params` vs.
   other collections in `variables` in mind. The original API was elegant because
   one only needed to pass around the optimizer, which included the parameters,

docs/notebooks/flax_basics.ipynb

@@ -382,7 +382,7 @@
         "\n",
         "@jax.jit\n",
         "def update_params(params, learning_rate, grads):\n",
-        "  params = jax.tree_multimap(\n",
+        "  params = jax.tree_map(\n",
         "      lambda p, g: p - learning_rate * g, params, grads)\n",
         "  return params\n",
         "\n",

docs/notebooks/jax_for_the_impatient.ipynb

@@ -1,18 +1,4 @@
 {
-  "nbformat": 4,
-  "nbformat_minor": 0,
-  "metadata": {
-    "colab": {
-      "name": "JAX for the impatient.ipynb",
-      "provenance": [],
-      "collapsed_sections": [],
-      "toc_visible": true
-    },
-    "kernelspec": {
-      "name": "python3",
-      "display_name": "Python 3"
-    }
-  },
   "cells": [
     {
       "cell_type": "markdown",
@@ -64,20 +50,14 @@
     },
     {
       "cell_type": "code",
+      "execution_count": 2,
       "metadata": {
-        "id": "L2HKiLTNJ4Eh",
-        "outputId": "c4297a1a-4e4b-4bdc-ca5d-3d33aca92b3b",
         "colab": {
           "base_uri": "https://localhost:8080/"
-        }
+        },
+        "id": "L2HKiLTNJ4Eh",
+        "outputId": "c4297a1a-4e4b-4bdc-ca5d-3d33aca92b3b"
       },
-      "source": [
-        "m = jnp.ones((4,4)) # We're generating one 4 by 4 matrix filled with ones.\n",
-        "n = jnp.array([[1.0, 2.0, 3.0, 4.0],\n",
-        "               [5.0, 6.0, 7.0, 8.0]]) # An explicit 2 by 4 array\n",
-        "m"
-      ],
-      "execution_count": 2,
       "outputs": [
         {
           "name": "stderr",
@@ -99,6 +79,12 @@
           "metadata": {},
           "output_type": "execute_result"
         }
+      ],
+      "source": [
+        "m = jnp.ones((4,4)) # We're generating one 4 by 4 matrix filled with ones.\n",
+        "n = jnp.array([[1.0, 2.0, 3.0, 4.0],\n",
+        "               [5.0, 6.0, 7.0, 8.0]]) # An explicit 2 by 4 array\n",
+        "m"
       ]
     },
     {
@@ -116,11 +102,11 @@
       "cell_type": "code",
       "execution_count": 3,
       "metadata": {
-        "id": "9do-ZRGaRThn",
-        "outputId": "9c4feb4d-3bd1-4921-97ce-c8087b37496f",
         "colab": {
           "base_uri": "https://localhost:8080/"
-        }
+        },
+        "id": "9do-ZRGaRThn",
+        "outputId": "9c4feb4d-3bd1-4921-97ce-c8087b37496f"
       },
       "outputs": [
         {
@@ -1069,7 +1055,7 @@
         "id": "3s167WGKUlZ9"
       },
       "source": [
-        "A more flexible version of `tree_map` would be `tree_multimap`. Instead of applying a standalone function to each of the tree leaves, you also provide a tuple of additional trees with similar shape to the input tree that will provide per leaf arguments to the function."
+        "Instead of applying a standalone function to each of the tree leaves, you can also provide a tuple of additional trees with similar shape to the input tree that will provide per leaf arguments to the function."
       ]
     },
     {
@@ -1095,7 +1081,8 @@
         }
       ],
       "source": [
-        "tree_util.tree_multimap(lambda x,y: x+y, t, tree_util.tree_map(lambda x: x*x, t))"
+        "t2 = tree_util.tree_map(lambda x: x*x, t)\n",
+        "tree_util.tree_map(lambda x,y: x+y, t, t2)"
       ]
     },
     {
@@ -1199,7 +1186,7 @@
         "id": "nW1IKnjqXFdN"
       },
       "source": [
-        "Now using our tree of params, we can write the gradient descent in a simpler way using `jax.tree_multimap`:"
+        "Now using our tree of params, we can write the gradient descent in a simpler way using `jax.tree_map`:"
       ]
     },
     {
@@ -1246,7 +1233,7 @@
         "# Always remember to jit!\n",
         "@jax.jit\n",
         "def update_params_pytree(params, learning_rate, x_samples, y_samples):\n",
-        "  params = jax.tree_multimap(\n",
+        "  params = jax.tree_map(\n",
         "        lambda p, g: p - learning_rate * g, params,\n",
         "        jax.grad(mse_pytree)(params, x_samples, y_samples))\n",
         "  return params\n",
@@ -1280,7 +1267,7 @@
         "for i in range(101):\n",
         "  # Note that here the loss is computed before the param update.\n",
         "    loss_val, grads = loss_grad_fn(params, x_samples, y_samples)\n",
-        "    params = jax.tree_multimap(\n",
+        "    params = jax.tree_map(\n",
         "        lambda p, g: p - learning_rate * g, params, grads)\n",
         "    if (i % 5 == 0):\n",
         "        print(f\"Loss step {i}: \", loss_val)"
@@ -1295,5 +1282,19 @@
         "That's all you needed to know to get started with Flax! To dive deeper, we very much recommend checking the JAX [docs](https://jax.readthedocs.io/en/latest/index.html)."
       ]
     }
-  ]
+  ],
+  "metadata": {
+    "colab": {
+      "collapsed_sections": [],
+      "name": "JAX for the impatient.ipynb",
+      "provenance": [],
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
 }

examples/imagenet/train.py

@@ -145,11 +145,11 @@ def loss_fn(params):
     # if is_fin == False the gradients contain Inf/NaNs and optimizer state and
     # params should be restored (= skip this step).
     new_state = new_state.replace(
-        opt_state=jax.tree_multimap(
+        opt_state=jax.tree_map(
             functools.partial(jnp.where, is_fin),
             new_state.opt_state,
             state.opt_state),
-        params=jax.tree_multimap(
+        params=jax.tree_map(
             functools.partial(jnp.where, is_fin),
             new_state.params,
             state.params))

examples/linen_design_test/linear_regression.py

@@ -45,4 +45,4 @@ def init_params(rng):
   loss, grad = jax.value_and_grad(loss_fn)(params)
   print(i, "loss = ", loss, "Yhat = ", predict(params))
   lr = 0.03
-  params = jax.tree_multimap(lambda x, d: x - lr * d, params, grad)
+  params = jax.tree_map(lambda x, d: x - lr * d, params, grad)

examples/wmt/train.py

@@ -224,9 +224,9 @@ def loss_fn(params):
     # params should be restored (= skip this step).
     select_fn = functools.partial(jnp.where, is_fin)
     new_state = new_state.replace(
-        opt_state=jax.tree_multimap(
+        opt_state=jax.tree_map(
             select_fn, new_state.opt_state, state.opt_state),
-        params=jax.tree_multimap(
+        params=jax.tree_map(
             select_fn, new_state.params, state.params)
         )
     metrics["loss_scale"] = dynamic_scale.scale * metrics["denominator"]

tests/core/core_frozen_dict_test.py

@@ -42,7 +42,7 @@ def test_frozen_dict_pop(self):
     self.assertEqual(unfreeze(b), {'b': {'c': 2}})
 
   def test_frozen_dict_partially_maps(self):
-    x = jax.tree_multimap(
+    x = jax.tree_map(
         lambda a, b: (a, b),
         freeze({'a': 2}), freeze({'a': {'b': 1}}))
     self.assertEqual(unfreeze(x), {'a': (2, {'b': 1})})

tests/linen/linen_module_test.py

@@ -44,7 +44,7 @@
 
 def tree_equals(x, y):
   return jax.tree_util.tree_all(
-      jax.tree_multimap(operator.eq, x, y))
+      jax.tree_map(operator.eq, x, y))
 
 
 class DummyModule(nn.Module):
@@ -1074,7 +1074,7 @@ def __call__(self, c, x):
       },
     })
     self.assertTrue(jax.tree_util.tree_all(
-        jax.tree_multimap(
+        jax.tree_map(
             lambda x, y: np.testing.assert_allclose(x, y, atol=1e-7),
             cntrs, ref_cntrs)
           ))