import tree_map from tree_util

docs/notebooks/07_Online_Time_Series_Prediction.ipynb

@@ -542,7 +542,7 @@
     }
    ],
    "source": [
-    "jax.tree_map(lambda x: x[-1], opt_info.params)"
+    "jax.tree_util.tree_map(lambda x: x[-1], opt_info.params)"
    ]
   },
   {
@@ -774,7 +774,7 @@
     "\n",
     "        def project_params(params: Any, opt_state: OptState = None):\n",
     "            del opt_state\n",
-    "            return jax.tree_map(lambda w: jnp.clip(w, -1, 1), params)\n",
+    "            return jax.tree_util.tree_map(lambda w: jnp.clip(w, -1, 1), params)\n",
     "\n",
     "        def params_predicate(m: str, n: str, p: jnp.ndarray) -> bool:\n",
     "            # print(m, n, p)\n",
@@ -1456,7 +1456,7 @@
     "    def fun_batch(*args, **kwargs):\n",
     "        res = VMap(fun)(*args, **kwargs)\n",
     "        if take_mean:\n",
-    "            res = jax.tree_map(lambda x: x.mean(axis=0), res)\n",
+    "            res = jax.tree_util.tree_map(lambda x: x.mean(axis=0), res)\n",
     "        return res\n",
     "\n",
     "    return fun_batch\n",
@@ -1681,7 +1681,7 @@
     "\n",
     "    BEST_STEP_SIZE[name] = loss.idxmin()\n",
     "    best_idx = loss.reset_index(drop=True).idxmin()\n",
-    "    BEST_GYM[name] = jax.tree_map(lambda x: x[:, best_idx], gym)\n",
+    "    BEST_GYM[name] = jax.tree_util.tree_map(lambda x: x[:, best_idx], gym)\n",
     "\n",
     "    ax = loss[loss < 0.15].plot(logx=True, logy=False, ax=ax, label=name)\n",
     "plt.legend()"
@@ -1972,7 +1972,7 @@
     "I_BEST_PARAM = jnp.argmin(x)\n",
     "\n",
     "\n",
-    "BEST_NEWTON_GYM = jax.tree_map(lambda x: x[:, I_BEST_PARAM], gym_newton)\n",
+    "BEST_NEWTON_GYM = jax.tree_util.tree_map(lambda x: x[:, I_BEST_PARAM], gym_newton)\n",
     "print(\"Best newton parameters: \", STEP_SIZE, NEWTON_EPS)"
    ]
   },

docs/notebooks/07_Online_Time_Series_Prediction.md

@@ -223,7 +223,7 @@ pd.Series(opt_info.loss).expanding().mean().plot()
 Let's look at the latest weights:
 
 ```python
-jax.tree_map(lambda x: x[-1], opt_info.params)
+jax.tree_util.tree_map(lambda x: x[-1], opt_info.params)
 ```
 
 
@@ -348,7 +348,7 @@ def build_agent(time_series_model=None, opt=None, embargo=1):
 
         def project_params(params: Any, opt_state: OptState = None):
             del opt_state
-            return jax.tree_map(lambda w: jnp.clip(w, -1, 1), params)
+            return jax.tree_util.tree_map(lambda w: jnp.clip(w, -1, 1), params)
 
         def params_predicate(m: str, n: str, p: jnp.ndarray) -> bool:
             # print(m, n, p)
@@ -563,7 +563,7 @@ def add_batch(fun, take_mean=True):
     def fun_batch(*args, **kwargs):
         res = VMap(fun)(*args, **kwargs)
         if take_mean:
-            res = jax.tree_map(lambda x: x.mean(axis=0), res)
+            res = jax.tree_util.tree_map(lambda x: x.mean(axis=0), res)
         return res
 
     return fun_batch
@@ -649,7 +649,7 @@ for name, (gym, info) in res.items():
 
     BEST_STEP_SIZE[name] = loss.idxmin()
     best_idx = loss.reset_index(drop=True).idxmin()
-    BEST_GYM[name] = jax.tree_map(lambda x: x[:, best_idx], gym)
+    BEST_GYM[name] = jax.tree_util.tree_map(lambda x: x[:, best_idx], gym)
 
     ax = loss[loss < 0.15].plot(logx=True, logy=False, ax=ax, label=name)
 plt.legend()
@@ -758,7 +758,7 @@ x = jnp.where(jnp.isnan(x), jnp.inf, x)
 I_BEST_PARAM = jnp.argmin(x)
 
 
-BEST_NEWTON_GYM = jax.tree_map(lambda x: x[:, I_BEST_PARAM], gym_newton)
+BEST_NEWTON_GYM = jax.tree_util.tree_map(lambda x: x[:, I_BEST_PARAM], gym_newton)
 print("Best newton parameters: ", STEP_SIZE, NEWTON_EPS)
 ```
 

docs/notebooks/07_Online_Time_Series_Prediction.py

@@ -208,7 +208,7 @@ def learn(y, X=None):
 
 # Let's look at the latest weights:
 
-jax.tree_map(lambda x: x[-1], opt_info.params)
+jax.tree_util.tree_map(lambda x: x[-1], opt_info.params)
 
 
 # ## Learn and Forecast
@@ -324,7 +324,7 @@ def model_with_loss(y, X=None):
 
         def project_params(params: Any, opt_state: OptState = None):
             del opt_state
-            return jax.tree_map(lambda w: jnp.clip(w, -1, 1), params)
+            return jax.tree_util.tree_map(lambda w: jnp.clip(w, -1, 1), params)
 
         def params_predicate(m: str, n: str, p: jnp.ndarray) -> bool:
             # print(m, n, p)
@@ -520,7 +520,7 @@ def add_batch(fun, take_mean=True):
     def fun_batch(*args, **kwargs):
         res = VMap(fun)(*args, **kwargs)
         if take_mean:
-            res = jax.tree_map(lambda x: x.mean(axis=0), res)
+            res = jax.tree_util.tree_map(lambda x: x.mean(axis=0), res)
         return res
 
     return fun_batch
@@ -602,7 +602,7 @@ def scan_params(step_size):
 
     BEST_STEP_SIZE[name] = loss.idxmin()
     best_idx = loss.reset_index(drop=True).idxmin()
-    BEST_GYM[name] = jax.tree_map(lambda x: x[:, best_idx], gym)
+    BEST_GYM[name] = jax.tree_util.tree_map(lambda x: x[:, best_idx], gym)
 
     ax = loss[loss < 0.15].plot(logx=True, logy=False, ax=ax, label=name)
 plt.legend()
@@ -696,7 +696,7 @@ def scan_params(hparams):
 I_BEST_PARAM = jnp.argmin(x)
 
 
-BEST_NEWTON_GYM = jax.tree_map(lambda x: x[:, I_BEST_PARAM], gym_newton)
+BEST_NEWTON_GYM = jax.tree_util.tree_map(lambda x: x[:, I_BEST_PARAM], gym_newton)
 print("Best newton parameters: ", STEP_SIZE, NEWTON_EPS)
 
 # +

docs/notebooks/08_Online_learning_in_non_stationary_environments.ipynb

@@ -151,7 +151,7 @@
     "\n",
     "        def project_params(params: Any, opt_state: OptState = None):\n",
     "            del opt_state\n",
-    "            return jax.tree_map(lambda w: jnp.clip(w, -1, 1), params)\n",
+    "            return jax.tree_util.tree_map(lambda w: jnp.clip(w, -1, 1), params)\n",
     "\n",
     "        def params_predicate(m: str, n: str, p: jnp.ndarray) -> bool:\n",
     "            # print(m, n, p)\n",
@@ -278,7 +278,7 @@
     "        BEST_STEP_SIZE[name] = loss.idxmin()\n",
     "\n",
     "        best_idx = jnp.argmax(gym.reward[LEARN_TIME_SLICE].mean(axis=0))\n",
-    "        BEST_GYM[name] = jax.tree_map(lambda x: x[:, best_idx], gym)\n",
+    "        BEST_GYM[name] = jax.tree_util.tree_map(lambda x: x[:, best_idx], gym)\n",
     "\n",
     "        ax = loss.plot(\n",
     "            logx=True, logy=False, ax=ax, label=name, ylim=(MIN_ERR, MAX_ERR)\n",
@@ -444,7 +444,7 @@
     "    x = jnp.where(jnp.isnan(x), jnp.inf, x)\n",
     "    I_BEST_PARAM = jnp.argmin(x)\n",
     "\n",
-    "    BEST_NEWTON_GYM = jax.tree_map(lambda x: x[:, I_BEST_PARAM], gym_newton)\n",
+    "    BEST_NEWTON_GYM = jax.tree_util.tree_map(lambda x: x[:, I_BEST_PARAM], gym_newton)\n",
     "\n",
     "    print(\"Best newton parameters: \", STEP_SIZE, NEWTON_EPS)\n",
     "\n",

docs/notebooks/08_Online_learning_in_non_stationary_environments.md

@@ -107,7 +107,7 @@ def build_agent(time_series_model=None, opt=None):
 
         def project_params(params: Any, opt_state: OptState = None):
             del opt_state
-            return jax.tree_map(lambda w: jnp.clip(w, -1, 1), params)
+            return jax.tree_util.tree_map(lambda w: jnp.clip(w, -1, 1), params)
 
         def params_predicate(m: str, n: str, p: jnp.ndarray) -> bool:
             # print(m, n, p)
@@ -203,7 +203,7 @@ def scan_hparams_first_order():
         BEST_STEP_SIZE[name] = loss.idxmin()
 
         best_idx = jnp.argmax(gym.reward[LEARN_TIME_SLICE].mean(axis=0))
-        BEST_GYM[name] = jax.tree_map(lambda x: x[:, best_idx], gym)
+        BEST_GYM[name] = jax.tree_util.tree_map(lambda x: x[:, best_idx], gym)
 
         ax = loss.plot(
             logx=True, logy=False, ax=ax, label=name, ylim=(MIN_ERR, MAX_ERR)
@@ -324,7 +324,7 @@ def scan_hparams_newton():
     x = jnp.where(jnp.isnan(x), jnp.inf, x)
     I_BEST_PARAM = jnp.argmin(x)
 
-    BEST_NEWTON_GYM = jax.tree_map(lambda x: x[:, I_BEST_PARAM], gym_newton)
+    BEST_NEWTON_GYM = jax.tree_util.tree_map(lambda x: x[:, I_BEST_PARAM], gym_newton)
 
     print("Best newton parameters: ", STEP_SIZE, NEWTON_EPS)
 

docs/notebooks/08_Online_learning_in_non_stationary_environments.py

@@ -102,7 +102,7 @@ def model_with_loss(y, X=None):
 
         def project_params(params: Any, opt_state: OptState = None):
             del opt_state
-            return jax.tree_map(lambda w: jnp.clip(w, -1, 1), params)
+            return jax.tree_util.tree_map(lambda w: jnp.clip(w, -1, 1), params)
 
         def params_predicate(m: str, n: str, p: jnp.ndarray) -> bool:
             # print(m, n, p)
@@ -195,7 +195,7 @@ def scan_params(step_size):
         BEST_STEP_SIZE[name] = loss.idxmin()
 
         best_idx = jnp.argmax(gym.reward[LEARN_TIME_SLICE].mean(axis=0))
-        BEST_GYM[name] = jax.tree_map(lambda x: x[:, best_idx], gym)
+        BEST_GYM[name] = jax.tree_util.tree_map(lambda x: x[:, best_idx], gym)
 
         ax = loss.plot(
             logx=True, logy=False, ax=ax, label=name, ylim=(MIN_ERR, MAX_ERR)
@@ -310,7 +310,7 @@ def scan_params(hparams):
     x = jnp.where(jnp.isnan(x), jnp.inf, x)
     I_BEST_PARAM = jnp.argmin(x)
 
-    BEST_NEWTON_GYM = jax.tree_map(lambda x: x[:, I_BEST_PARAM], gym_newton)
+    BEST_NEWTON_GYM = jax.tree_util.tree_map(lambda x: x[:, I_BEST_PARAM], gym_newton)
 
     print("Best newton parameters: ", STEP_SIZE, NEWTON_EPS)
 

wax/modules/func_optimizer.py

@@ -72,7 +72,7 @@ def __call__(self, *args, **kwargs):
         if self.grads_fill_nan_inf:
             grads = FillNanInf()(grads)
 
-        # trainable_params = jax.tree_map(self.opt, trainable_params, grads)
+        # trainable_params = jax.tree_util.tree_map(self.opt, trainable_params, grads)
         trainable_params = self.opt(trainable_params, grads)
 
         hk.set_state("trainable_params", trainable_params)

wax/modules/snarimax_test.py

@@ -119,7 +119,7 @@ def model_with_loss(y, X=None):
 
         def project_params(params: Any, opt_state: OptState = None):
             del opt_state
-            return jax.tree_map(lambda w: jnp.clip(w, -1, 1), params)
+            return jax.tree_util.tree_map(lambda w: jnp.clip(w, -1, 1), params)
 
         def params_predicate(m: str, n: str, p: jnp.ndarray) -> bool:
             # print(m, n, p)

wax/modules/vmap.py

@@ -51,7 +51,7 @@ def fun_batch(*args, **kwargs):
             res = vmap_lift_with_state(fun, split_rng=False)(*args, **kwargs)
 
         if take_mean:
-            res = jax.tree_map(lambda x: x.mean(axis=0), res)
+            res = jax.tree_util.tree_map(lambda x: x.mean(axis=0), res)
         return res
 
     return fun_batch

wax/optim/newton.py

@@ -75,7 +75,7 @@ def scale_by_newton(eps: float = 1e-7) -> base.GradientTransformation:
     """
 
     def init_fn(params):
-        hessian_inv = jax.tree_map(
+        hessian_inv = jax.tree_util.tree_map(
             lambda t: jnp.eye(len(t.flatten()), dtype=t.dtype) / eps, params
         )
         return ScaleByNewtonState(hessian_inv=hessian_inv)
@@ -89,13 +89,15 @@ class Tuple(tuple):
 
             ...
 
-        shapes = jax.tree_map(lambda x: Tuple(x.shape), updates)
-        updates = jax.tree_map(lambda x: x.flatten(), updates)
-        hessian_inv = jax.tree_map(
+        shapes = jax.tree_util.tree_map(lambda x: Tuple(x.shape), updates)
+        updates = jax.tree_util.tree_map(lambda x: x.flatten(), updates)
+        hessian_inv = jax.tree_util.tree_map(
             lambda u, hinv: sherman_morrison(hinv, u, u), updates, state.hessian_inv
         )
-        updates = jax.tree_map(lambda hinv, g: hinv @ g, hessian_inv, updates)
-        updates = jax.tree_map(lambda u, shape: u.reshape(shape), updates, shapes)
+        updates = jax.tree_util.tree_map(lambda hinv, g: hinv @ g, hessian_inv, updates)
+        updates = jax.tree_util.tree_map(
+            lambda u, shape: u.reshape(shape), updates, shapes
+        )
 
         return updates, ScaleByNewtonState(hessian_inv=hessian_inv)