We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Optimizer tree_init returns a slots list, but tree_update returns a slots tuple
### Description

`tree_init` and `tree_update` are not consistent. One returns a list for the slots, the other a tuple.

It is a super minor detail but I was trying to conditionally run a `tree_update` with `jax.cond`, and this minor difference made that break, since the PyTreeDefs were different.

Casting the slots list comprehension to a tuple  ([here](https://github.com/google/trax/blob/0ca17db895c7d9bb203e66e074f49e9481b87513/trax/optimizers/base.py#L119-L120)) solved this for me, but I'm not sure if you want to go with tuple or list so I raise an issue instead of PR.

### Environment information
```
OS: Ubuntu 18.04

$ pip freeze | grep trax
-e git+git@github.com:google/trax.git@0ca17db895c7d9bb203e66e074f49e9481b87513#egg=trax
(latest commit from Sep 30)

$ pip freeze | grep tensor
tensorflow==2.3.1

$ pip freeze | grep jax
jax==0.2.0
jaxlib @ https://storage.googleapis.com/jax-releases/cuda110/jaxlib-0.1.55-cp36-none-manylinux2010_x86_64.whl

$ python -V
Python 3.6.9
```

### For bugs: reproduction and error logs

You can add the following lines to `optimizers_test.py` and see the behavior.

```
# Steps to reproduce:


    # Show that tree_update returns slots in a tuple not list
    old_slots = opt_2.slots
    grad_tree = np.zeros_like(weight_tree)
    _, new_slots, _ = opt_2.tree_update(1, grad_tree, weight_tree, opt_2.slots, opt_2.opt_params)
    self.assertIsInstance(old_slots, list)  # PASS
    self.assertIsInstance(opt_2.slots, list)  # FAIL. it's a tuple
    self.assertIsInstance(new_slots, list)  # FAIL. it's a tuple

```

```
# Error logs:


TypeError: true_fun and false_fun output must have same type structure, got PyTreeDef(tuple, [PyTreeDef(dict[['dyn']], [PyTreeDef(dict[['ff']], [PyTreeDef(dict[['dense0', 'dense1', 'dense2', 'dense_last']], [PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*])])])]),PyTreeDef(tuple, [PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*])]),PyTreeDef(dict[['gradients_l2', 'weights_l2']], [*,*])]) and PyTreeDef(tuple, [PyTreeDef(dict[['dyn']], [PyTreeDef(dict[['ff']], [PyTreeDef(dict[['dense0', 'dense1', 'dense2', 'dense_last']], [PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*])])])]),PyTreeDef(list, [PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*])]),PyTreeDef(dict[['gradients_l2', 'weights_l2']], [*,*])]).

```


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `trax/optimizers/base.py`
Content:
```
1 # coding=utf-8
2 # Copyright 2020 The Trax Authors.
3 #
4 # Licensed under the Apache License, Version 2.0 (the "License");
5 # you may not use this file except in compliance with the License.
6 # You may obtain a copy of the License at
7 #
8 #     http://www.apache.org/licenses/LICENSE-2.0
9 #
10 # Unless required by applicable law or agreed to in writing, software
11 # distributed under the License is distributed on an "AS IS" BASIS,
12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 # See the License for the specific language governing permissions and
14 # limitations under the License.
15 
16 # Lint as: python3
17 """Trax base optimizer class."""
18 
19 from trax import fastmath
20 from trax.fastmath import numpy as jnp
21 
22 
23 class Optimizer(object):
24   """Base class for optimizers that work hand in hand with Trax layers.
25 
26   To define an optimizer subclass, specify its behavior with respect to a
27   single level/node in the network (e.g., a single dense layer):
28 
29     - `init`: how to create/initialize optimizer-internal weights ("slots")
30         whose shape matches the node's weight shape.
31     - `update`: how to use gradient information to update node weights and
32         optimizer slots.
33 
34   The Trax runtime combines these node-local computations into weight updates
35   and slot updates for the whole tree of layers in the model.
36   """
37 
38   def __init__(self, learning_rate=0.01, clip_grad_norm=None,
39                **init_opt_params):
40     """Sets initial hyperparameter values for this optimizer.
41 
42     Takes initial optimizer parameters as keyword arguments. These values can
43     be changed between training steps, e.g., for learning rate schedules.
44 
45     If you want your subclass to expose hyperparameters for gin configuration,
46     override this constructor and use explicitly named keyword arguments. See
47     `momentum.Momentum.__init__` for one such example.
48 
49     Args:
50       learning_rate: The initial learning rate.
51       clip_grad_norm: float; the value to which gradients will be clipped.
52       **init_opt_params: Initial values of any additional optimizer parameters.
53     """
54     init_opt_params['learning_rate'] = learning_rate
55     self._init_opt_params = {
56         name: jnp.array(value) for (name, value) in init_opt_params.items()
57     }
58     self._slots = None
59     # Gradient clipping happens with respect to the norm of the whole gradient
60     # tree, so it is not passed to single-slot updates, but done in this class
61     # for the whole gradient tree.
62     self._clip_grad_norm = clip_grad_norm
63 
64   def init(self, weights):
65     """Creates optimizer slots for the given parameters.
66 
67     Args:
68       weights: Trainable weights for one layer. Optimizer slots typically match
69           the data shape and type of the given layer weights.
70     """
71     raise NotImplementedError
72 
73   def update(self, step, grads, weights, slots, opt_params):
74     """Computes one step's worth of updates.
75 
76     The update computes both new weights for the layer/node and new slot values
77     for the optimizer.
78 
79     Args:
80       step: Current step number in the training process.
81       grads: Gradients for the weights of the sublayer.
82       weights: Current weights for the sublayer.
83       slots: Optimizer slots.
84       opt_params: Optimizer hyperparameters (e.g. learning rate, momentum).
85 
86     Returns:
87       Tuple of (new_weights, new_slots).
88     """
89     raise NotImplementedError
90 
91   @property
92   def slots(self):
93     return self._slots
94 
95   @slots.setter
96   def slots(self, slots):
97     self._slots = slots
98 
99   @property
100   def opt_params(self):
101     return self._init_opt_params
102 
103   @opt_params.setter
104   def opt_params(self, opt_params):
105     self._init_opt_params = opt_params
106 
107   def tree_init(self, weight_tree):
108     """Assembles node-local initializations into full-tree initialization.
109 
110     Args:
111       weight_tree: Weights for an entire model, in a tree that matches the
112           model's layer structure.
113 
114     Returns:
115       Tuple `(slots, opt_params)`, where `slots` are the initialized optimizer
116       slot values and `opt_params` are optimizer hyperparameters (e.g.,
117       learning rate, momentum).
118     """
119     self._slots = [self.init(weight)
120                    for weight in fastmath.tree_flatten(weight_tree)]
121     return (
122         self._slots,
123         self._init_opt_params,
124     )
125 
126   def tree_update(self, step, grad_tree, weight_tree, slots, opt_params):
127     """Assembles node-local weight and slot updates for the full layer tree.
128 
129     Args:
130       step: Current step number in the training process.
131       grad_tree: Gradients for the entire model, in a tree that matches the
132           model's layer structure.
133       weight_tree: Current weights for the entire model, in a tree that matches
134           the model's layer structure.
135       slots: Optimizer slots.
136       opt_params: Optimizer hyperparameters (e.g. learning rate, momentum).
137 
138     Returns:
139       Tuple `(weights, slots)`, where `weights` are the optimizer-updated
140       weights for the whole model (in a tree matching the model's layer
141       structure) and `slots` are the updated optimizer slot values.
142     """
143     grads_flat = fastmath.tree_flatten(grad_tree)
144     grads_norm = self._l2_norm(grads_flat)
145     if self._clip_grad_norm is not None:
146       max_norm = self._clip_grad_norm
147       grads_flat = [jnp.where(grads_norm < max_norm,  # pylint: disable=g-complex-comprehension
148                               g,
149                               g * (max_norm / grads_norm))
150                     for g in grads_flat]
151     weights_flat = fastmath.tree_flatten(weight_tree)
152     weights_norm = self._l2_norm(weights_flat)
153     updated_pairs = [
154         self._update_and_check(step, grad, weight, slot, opt_params)
155         for (grad, weight, slot) in zip(grads_flat, weights_flat, slots)
156     ]
157     new_weights_flat, self.slots = zip(*updated_pairs)
158     new_weights, _ = fastmath.tree_unflatten(new_weights_flat, weight_tree)
159     metrics = {'gradients_l2': grads_norm, 'weights_l2': weights_norm}
160     return new_weights, self.slots, metrics
161 
162   def _l2_norm(self, flat_list):
163     """Returns the aggregate L2 norm of a list of tensors."""
164     if fastmath.is_backend(fastmath.Backend.JAX):
165       norm = jnp.sqrt(sum(jnp.vdot(x, x) for x in flat_list))
166     else:  # TODO(lukaszkaiser): add vdot to TF-numpy
167       norm = jnp.sqrt(sum(jnp.sum(x*x) for x in flat_list))
168     return norm
169 
170   def _update_and_check(self, step, grads, weights, slots, opt_params):
171     """Updates a single weight array and checks types."""
172     new_weights, new_slots = self.update(
173         step, grads, weights, slots, opt_params)
174     if isinstance(weights, jnp.ndarray):
175       if not isinstance(new_weights, jnp.ndarray):
176         raise ValueError(
177             f'New weight values should be of type jnp.ndarray or a subclass; '
178             f'instead got {type(new_weights)}.')
179       if new_weights.dtype != weights.dtype:
180         raise ValueError(
181             f'New weight values dtype ({new_weights.dtype}) does not match '
182             f'the old one ({weights.dtype}).')
183     return new_weights, new_slots
184 
185 
186 class SGD(Optimizer):
187   """Stochastic gradient descent (SGD) optimizer.
188 
189   A simple optimizer with no weights ("slots") of its own.
190   """
191 
192   def init(self, weights):
193     return None
194 
195   def update(self, step, grads, weights, slots, opt_params):
196     del step, slots
197     lr = opt_params['learning_rate']
198     new_weights = weights - (lr * grads).astype(weights.dtype)
199     return new_weights, None
200 
201 
202 # Utilities.
203 
204 
205 def l2_norm(tree):
206   """Compute the l2 norm of a pytree of arrays. Useful for weight decay."""
207   leaves = fastmath.tree_flatten(tree)
208   return jnp.sqrt(sum(jnp.vdot(x, x) for x in leaves))
209 
210 
211 def clip_grads(grad_tree, max_norm):
212   """Clip gradients stored as a pytree of arrays to maximum norm `max_norm`."""
213   norm = l2_norm(grad_tree)
214   normalize = lambda g: jnp.where(norm < max_norm, g, g * (max_norm / norm))
215   return fastmath.nested_map(grad_tree, normalize)
216
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/trax/optimizers/base.py b/trax/optimizers/base.py
--- a/trax/optimizers/base.py
+++ b/trax/optimizers/base.py
@@ -154,7 +154,7 @@
         self._update_and_check(step, grad, weight, slot, opt_params)
         for (grad, weight, slot) in zip(grads_flat, weights_flat, slots)
     ]
-    new_weights_flat, self.slots = zip(*updated_pairs)
+    new_weights_flat, self.slots = map(list, zip(*updated_pairs))
     new_weights, _ = fastmath.tree_unflatten(new_weights_flat, weight_tree)
     metrics = {'gradients_l2': grads_norm, 'weights_l2': weights_norm}
     return new_weights, self.slots, metrics

{"golden_diff": "diff --git a/trax/optimizers/base.py b/trax/optimizers/base.py\n--- a/trax/optimizers/base.py\n+++ b/trax/optimizers/base.py\n@@ -154,7 +154,7 @@\n         self._update_and_check(step, grad, weight, slot, opt_params)\n         for (grad, weight, slot) in zip(grads_flat, weights_flat, slots)\n     ]\n-    new_weights_flat, self.slots = zip(*updated_pairs)\n+    new_weights_flat, self.slots = map(list, zip(*updated_pairs))\n     new_weights, _ = fastmath.tree_unflatten(new_weights_flat, weight_tree)\n     metrics = {'gradients_l2': grads_norm, 'weights_l2': weights_norm}\n     return new_weights, self.slots, metrics\n", "issue": "Optimizer tree_init returns a slots list, but tree_update returns a slots tuple\n### Description\r\n\r\n`tree_init` and `tree_update` are not consistent. One returns a list for the slots, the other a tuple.\r\n\r\nIt is a super minor detail but I was trying to conditionally run a `tree_update` with `jax.cond`, and this minor difference made that break, since the PyTreeDefs were different.\r\n\r\nCasting the slots list comprehension to a tuple  ([here](https://github.com/google/trax/blob/0ca17db895c7d9bb203e66e074f49e9481b87513/trax/optimizers/base.py#L119-L120)) solved this for me, but I'm not sure if you want to go with tuple or list so I raise an issue instead of PR.\r\n\r\n### Environment information\r\n```\r\nOS: Ubuntu 18.04\r\n\r\n$ pip freeze | grep trax\r\n-e git+git@github.com:google/trax.git@0ca17db895c7d9bb203e66e074f49e9481b87513#egg=trax\r\n(latest commit from Sep 30)\r\n\r\n$ pip freeze | grep tensor\r\ntensorflow==2.3.1\r\n\r\n$ pip freeze | grep jax\r\njax==0.2.0\r\njaxlib @ https://storage.googleapis.com/jax-releases/cuda110/jaxlib-0.1.55-cp36-none-manylinux2010_x86_64.whl\r\n\r\n$ python -V\r\nPython 3.6.9\r\n```\r\n\r\n### For bugs: reproduction and error logs\r\n\r\nYou can add the following lines to `optimizers_test.py` and see the behavior.\r\n\r\n```\r\n# Steps to reproduce:\r\n\r\n\r\n    # Show that tree_update returns slots in a tuple not list\r\n    old_slots = opt_2.slots\r\n    grad_tree = np.zeros_like(weight_tree)\r\n    _, new_slots, _ = opt_2.tree_update(1, grad_tree, weight_tree, opt_2.slots, opt_2.opt_params)\r\n    self.assertIsInstance(old_slots, list)  # PASS\r\n    self.assertIsInstance(opt_2.slots, list)  # FAIL. it's a tuple\r\n    self.assertIsInstance(new_slots, list)  # FAIL. it's a tuple\r\n\r\n```\r\n\r\n```\r\n# Error logs:\r\n\r\n\r\nTypeError: true_fun and false_fun output must have same type structure, got PyTreeDef(tuple, [PyTreeDef(dict[['dyn']], [PyTreeDef(dict[['ff']], [PyTreeDef(dict[['dense0', 'dense1', 'dense2', 'dense_last']], [PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*])])])]),PyTreeDef(tuple, [PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*])]),PyTreeDef(dict[['gradients_l2', 'weights_l2']], [*,*])]) and PyTreeDef(tuple, [PyTreeDef(dict[['dyn']], [PyTreeDef(dict[['ff']], [PyTreeDef(dict[['dense0', 'dense1', 'dense2', 'dense_last']], [PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*])])])]),PyTreeDef(list, [PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*]),PyTreeDef(tuple, [*,*])]),PyTreeDef(dict[['gradients_l2', 'weights_l2']], [*,*])]).\r\n\r\n```\r\n\n", "before_files": [{"content": "# coding=utf-8\n# Copyright 2020 The Trax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# Lint as: python3\n\"\"\"Trax base optimizer class.\"\"\"\n\nfrom trax import fastmath\nfrom trax.fastmath import numpy as jnp\n\n\nclass Optimizer(object):\n  \"\"\"Base class for optimizers that work hand in hand with Trax layers.\n\n  To define an optimizer subclass, specify its behavior with respect to a\n  single level/node in the network (e.g., a single dense layer):\n\n    - `init`: how to create/initialize optimizer-internal weights (\"slots\")\n        whose shape matches the node's weight shape.\n    - `update`: how to use gradient information to update node weights and\n        optimizer slots.\n\n  The Trax runtime combines these node-local computations into weight updates\n  and slot updates for the whole tree of layers in the model.\n  \"\"\"\n\n  def __init__(self, learning_rate=0.01, clip_grad_norm=None,\n               **init_opt_params):\n    \"\"\"Sets initial hyperparameter values for this optimizer.\n\n    Takes initial optimizer parameters as keyword arguments. These values can\n    be changed between training steps, e.g., for learning rate schedules.\n\n    If you want your subclass to expose hyperparameters for gin configuration,\n    override this constructor and use explicitly named keyword arguments. See\n    `momentum.Momentum.__init__` for one such example.\n\n    Args:\n      learning_rate: The initial learning rate.\n      clip_grad_norm: float; the value to which gradients will be clipped.\n      **init_opt_params: Initial values of any additional optimizer parameters.\n    \"\"\"\n    init_opt_params['learning_rate'] = learning_rate\n    self._init_opt_params = {\n        name: jnp.array(value) for (name, value) in init_opt_params.items()\n    }\n    self._slots = None\n    # Gradient clipping happens with respect to the norm of the whole gradient\n    # tree, so it is not passed to single-slot updates, but done in this class\n    # for the whole gradient tree.\n    self._clip_grad_norm = clip_grad_norm\n\n  def init(self, weights):\n    \"\"\"Creates optimizer slots for the given parameters.\n\n    Args:\n      weights: Trainable weights for one layer. Optimizer slots typically match\n          the data shape and type of the given layer weights.\n    \"\"\"\n    raise NotImplementedError\n\n  def update(self, step, grads, weights, slots, opt_params):\n    \"\"\"Computes one step's worth of updates.\n\n    The update computes both new weights for the layer/node and new slot values\n    for the optimizer.\n\n    Args:\n      step: Current step number in the training process.\n      grads: Gradients for the weights of the sublayer.\n      weights: Current weights for the sublayer.\n      slots: Optimizer slots.\n      opt_params: Optimizer hyperparameters (e.g. learning rate, momentum).\n\n    Returns:\n      Tuple of (new_weights, new_slots).\n    \"\"\"\n    raise NotImplementedError\n\n  @property\n  def slots(self):\n    return self._slots\n\n  @slots.setter\n  def slots(self, slots):\n    self._slots = slots\n\n  @property\n  def opt_params(self):\n    return self._init_opt_params\n\n  @opt_params.setter\n  def opt_params(self, opt_params):\n    self._init_opt_params = opt_params\n\n  def tree_init(self, weight_tree):\n    \"\"\"Assembles node-local initializations into full-tree initialization.\n\n    Args:\n      weight_tree: Weights for an entire model, in a tree that matches the\n          model's layer structure.\n\n    Returns:\n      Tuple `(slots, opt_params)`, where `slots` are the initialized optimizer\n      slot values and `opt_params` are optimizer hyperparameters (e.g.,\n      learning rate, momentum).\n    \"\"\"\n    self._slots = [self.init(weight)\n                   for weight in fastmath.tree_flatten(weight_tree)]\n    return (\n        self._slots,\n        self._init_opt_params,\n    )\n\n  def tree_update(self, step, grad_tree, weight_tree, slots, opt_params):\n    \"\"\"Assembles node-local weight and slot updates for the full layer tree.\n\n    Args:\n      step: Current step number in the training process.\n      grad_tree: Gradients for the entire model, in a tree that matches the\n          model's layer structure.\n      weight_tree: Current weights for the entire model, in a tree that matches\n          the model's layer structure.\n      slots: Optimizer slots.\n      opt_params: Optimizer hyperparameters (e.g. learning rate, momentum).\n\n    Returns:\n      Tuple `(weights, slots)`, where `weights` are the optimizer-updated\n      weights for the whole model (in a tree matching the model's layer\n      structure) and `slots` are the updated optimizer slot values.\n    \"\"\"\n    grads_flat = fastmath.tree_flatten(grad_tree)\n    grads_norm = self._l2_norm(grads_flat)\n    if self._clip_grad_norm is not None:\n      max_norm = self._clip_grad_norm\n      grads_flat = [jnp.where(grads_norm < max_norm,  # pylint: disable=g-complex-comprehension\n                              g,\n                              g * (max_norm / grads_norm))\n                    for g in grads_flat]\n    weights_flat = fastmath.tree_flatten(weight_tree)\n    weights_norm = self._l2_norm(weights_flat)\n    updated_pairs = [\n        self._update_and_check(step, grad, weight, slot, opt_params)\n        for (grad, weight, slot) in zip(grads_flat, weights_flat, slots)\n    ]\n    new_weights_flat, self.slots = zip(*updated_pairs)\n    new_weights, _ = fastmath.tree_unflatten(new_weights_flat, weight_tree)\n    metrics = {'gradients_l2': grads_norm, 'weights_l2': weights_norm}\n    return new_weights, self.slots, metrics\n\n  def _l2_norm(self, flat_list):\n    \"\"\"Returns the aggregate L2 norm of a list of tensors.\"\"\"\n    if fastmath.is_backend(fastmath.Backend.JAX):\n      norm = jnp.sqrt(sum(jnp.vdot(x, x) for x in flat_list))\n    else:  # TODO(lukaszkaiser): add vdot to TF-numpy\n      norm = jnp.sqrt(sum(jnp.sum(x*x) for x in flat_list))\n    return norm\n\n  def _update_and_check(self, step, grads, weights, slots, opt_params):\n    \"\"\"Updates a single weight array and checks types.\"\"\"\n    new_weights, new_slots = self.update(\n        step, grads, weights, slots, opt_params)\n    if isinstance(weights, jnp.ndarray):\n      if not isinstance(new_weights, jnp.ndarray):\n        raise ValueError(\n            f'New weight values should be of type jnp.ndarray or a subclass; '\n            f'instead got {type(new_weights)}.')\n      if new_weights.dtype != weights.dtype:\n        raise ValueError(\n            f'New weight values dtype ({new_weights.dtype}) does not match '\n            f'the old one ({weights.dtype}).')\n    return new_weights, new_slots\n\n\nclass SGD(Optimizer):\n  \"\"\"Stochastic gradient descent (SGD) optimizer.\n\n  A simple optimizer with no weights (\"slots\") of its own.\n  \"\"\"\n\n  def init(self, weights):\n    return None\n\n  def update(self, step, grads, weights, slots, opt_params):\n    del step, slots\n    lr = opt_params['learning_rate']\n    new_weights = weights - (lr * grads).astype(weights.dtype)\n    return new_weights, None\n\n\n# Utilities.\n\n\ndef l2_norm(tree):\n  \"\"\"Compute the l2 norm of a pytree of arrays. Useful for weight decay.\"\"\"\n  leaves = fastmath.tree_flatten(tree)\n  return jnp.sqrt(sum(jnp.vdot(x, x) for x in leaves))\n\n\ndef clip_grads(grad_tree, max_norm):\n  \"\"\"Clip gradients stored as a pytree of arrays to maximum norm `max_norm`.\"\"\"\n  norm = l2_norm(grad_tree)\n  normalize = lambda g: jnp.where(norm < max_norm, g, g * (max_norm / norm))\n  return fastmath.nested_map(grad_tree, normalize)\n", "path": "trax/optimizers/base.py"}], "after_files": [{"content": "# coding=utf-8\n# Copyright 2020 The Trax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# Lint as: python3\n\"\"\"Trax base optimizer class.\"\"\"\n\nfrom trax import fastmath\nfrom trax.fastmath import numpy as jnp\n\n\nclass Optimizer(object):\n  \"\"\"Base class for optimizers that work hand in hand with Trax layers.\n\n  To define an optimizer subclass, specify its behavior with respect to a\n  single level/node in the network (e.g., a single dense layer):\n\n    - `init`: how to create/initialize optimizer-internal weights (\"slots\")\n        whose shape matches the node's weight shape.\n    - `update`: how to use gradient information to update node weights and\n        optimizer slots.\n\n  The Trax runtime combines these node-local computations into weight updates\n  and slot updates for the whole tree of layers in the model.\n  \"\"\"\n\n  def __init__(self, learning_rate=0.01, clip_grad_norm=None,\n               **init_opt_params):\n    \"\"\"Sets initial hyperparameter values for this optimizer.\n\n    Takes initial optimizer parameters as keyword arguments. These values can\n    be changed between training steps, e.g., for learning rate schedules.\n\n    If you want your subclass to expose hyperparameters for gin configuration,\n    override this constructor and use explicitly named keyword arguments. See\n    `momentum.Momentum.__init__` for one such example.\n\n    Args:\n      learning_rate: The initial learning rate.\n      clip_grad_norm: float; the value to which gradients will be clipped.\n      **init_opt_params: Initial values of any additional optimizer parameters.\n    \"\"\"\n    init_opt_params['learning_rate'] = learning_rate\n    self._init_opt_params = {\n        name: jnp.array(value) for (name, value) in init_opt_params.items()\n    }\n    self._slots = None\n    # Gradient clipping happens with respect to the norm of the whole gradient\n    # tree, so it is not passed to single-slot updates, but done in this class\n    # for the whole gradient tree.\n    self._clip_grad_norm = clip_grad_norm\n\n  def init(self, weights):\n    \"\"\"Creates optimizer slots for the given parameters.\n\n    Args:\n      weights: Trainable weights for one layer. Optimizer slots typically match\n          the data shape and type of the given layer weights.\n    \"\"\"\n    raise NotImplementedError\n\n  def update(self, step, grads, weights, slots, opt_params):\n    \"\"\"Computes one step's worth of updates.\n\n    The update computes both new weights for the layer/node and new slot values\n    for the optimizer.\n\n    Args:\n      step: Current step number in the training process.\n      grads: Gradients for the weights of the sublayer.\n      weights: Current weights for the sublayer.\n      slots: Optimizer slots.\n      opt_params: Optimizer hyperparameters (e.g. learning rate, momentum).\n\n    Returns:\n      Tuple of (new_weights, new_slots).\n    \"\"\"\n    raise NotImplementedError\n\n  @property\n  def slots(self):\n    return self._slots\n\n  @slots.setter\n  def slots(self, slots):\n    self._slots = slots\n\n  @property\n  def opt_params(self):\n    return self._init_opt_params\n\n  @opt_params.setter\n  def opt_params(self, opt_params):\n    self._init_opt_params = opt_params\n\n  def tree_init(self, weight_tree):\n    \"\"\"Assembles node-local initializations into full-tree initialization.\n\n    Args:\n      weight_tree: Weights for an entire model, in a tree that matches the\n          model's layer structure.\n\n    Returns:\n      Tuple `(slots, opt_params)`, where `slots` are the initialized optimizer\n      slot values and `opt_params` are optimizer hyperparameters (e.g.,\n      learning rate, momentum).\n    \"\"\"\n    self._slots = [self.init(weight)\n                   for weight in fastmath.tree_flatten(weight_tree)]\n    return (\n        self._slots,\n        self._init_opt_params,\n    )\n\n  def tree_update(self, step, grad_tree, weight_tree, slots, opt_params):\n    \"\"\"Assembles node-local weight and slot updates for the full layer tree.\n\n    Args:\n      step: Current step number in the training process.\n      grad_tree: Gradients for the entire model, in a tree that matches the\n          model's layer structure.\n      weight_tree: Current weights for the entire model, in a tree that matches\n          the model's layer structure.\n      slots: Optimizer slots.\n      opt_params: Optimizer hyperparameters (e.g. learning rate, momentum).\n\n    Returns:\n      Tuple `(weights, slots)`, where `weights` are the optimizer-updated\n      weights for the whole model (in a tree matching the model's layer\n      structure) and `slots` are the updated optimizer slot values.\n    \"\"\"\n    grads_flat = fastmath.tree_flatten(grad_tree)\n    grads_norm = self._l2_norm(grads_flat)\n    if self._clip_grad_norm is not None:\n      max_norm = self._clip_grad_norm\n      grads_flat = [jnp.where(grads_norm < max_norm,  # pylint: disable=g-complex-comprehension\n                              g,\n                              g * (max_norm / grads_norm))\n                    for g in grads_flat]\n    weights_flat = fastmath.tree_flatten(weight_tree)\n    weights_norm = self._l2_norm(weights_flat)\n    updated_pairs = [\n        self._update_and_check(step, grad, weight, slot, opt_params)\n        for (grad, weight, slot) in zip(grads_flat, weights_flat, slots)\n    ]\n    new_weights_flat, self.slots = map(list, zip(*updated_pairs))\n    new_weights, _ = fastmath.tree_unflatten(new_weights_flat, weight_tree)\n    metrics = {'gradients_l2': grads_norm, 'weights_l2': weights_norm}\n    return new_weights, self.slots, metrics\n\n  def _l2_norm(self, flat_list):\n    \"\"\"Returns the aggregate L2 norm of a list of tensors.\"\"\"\n    if fastmath.is_backend(fastmath.Backend.JAX):\n      norm = jnp.sqrt(sum(jnp.vdot(x, x) for x in flat_list))\n    else:  # TODO(lukaszkaiser): add vdot to TF-numpy\n      norm = jnp.sqrt(sum(jnp.sum(x*x) for x in flat_list))\n    return norm\n\n  def _update_and_check(self, step, grads, weights, slots, opt_params):\n    \"\"\"Updates a single weight array and checks types.\"\"\"\n    new_weights, new_slots = self.update(\n        step, grads, weights, slots, opt_params)\n    if isinstance(weights, jnp.ndarray):\n      if not isinstance(new_weights, jnp.ndarray):\n        raise ValueError(\n            f'New weight values should be of type jnp.ndarray or a subclass; '\n            f'instead got {type(new_weights)}.')\n      if new_weights.dtype != weights.dtype:\n        raise ValueError(\n            f'New weight values dtype ({new_weights.dtype}) does not match '\n            f'the old one ({weights.dtype}).')\n    return new_weights, new_slots\n\n\nclass SGD(Optimizer):\n  \"\"\"Stochastic gradient descent (SGD) optimizer.\n\n  A simple optimizer with no weights (\"slots\") of its own.\n  \"\"\"\n\n  def init(self, weights):\n    return None\n\n  def update(self, step, grads, weights, slots, opt_params):\n    del step, slots\n    lr = opt_params['learning_rate']\n    new_weights = weights - (lr * grads).astype(weights.dtype)\n    return new_weights, None\n\n\n# Utilities.\n\n\ndef l2_norm(tree):\n  \"\"\"Compute the l2 norm of a pytree of arrays. Useful for weight decay.\"\"\"\n  leaves = fastmath.tree_flatten(tree)\n  return jnp.sqrt(sum(jnp.vdot(x, x) for x in leaves))\n\n\ndef clip_grads(grad_tree, max_norm):\n  \"\"\"Clip gradients stored as a pytree of arrays to maximum norm `max_norm`.\"\"\"\n  norm = l2_norm(grad_tree)\n  normalize = lambda g: jnp.where(norm < max_norm, g, g * (max_norm / norm))\n  return fastmath.nested_map(grad_tree, normalize)\n", "path": "trax/optimizers/base.py"}]}