Repository: google/trax

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']], [*,*])]).

```



Discussion:

Resolving patch:
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 @@ def tree_update(self, step, grad_tree, weight_tree, slots, opt_params):
         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


Test patch: