Repository: google/flax

Issue:
Deserialized TrainState dosn't pass `_check_tree_and_avals` check triggered by jax control flow
### Problem you have encountered:
Training fails to run with restored `TrainState` when `jax.lax.cond` is in the loop. Specifically the `true_fun` and `false_fun` return type check fails because of having `optax._src...<TypeName>` vs `flax.serialization.<TypeName>` for updated and not-yet-updated branches respectively. 

### What you expected to happen:
Resume training without having any issues.

### Error message:
```
UnfilteredStackTrace: TypeError: true_fun and false_fun output must have same type structure, got PyTreeDef(CustomNode(<class 'flax.training.train_state.TrainState'>[(<bound method Module.apply of MLP(
    # attributes
    features = [4, 1]
)>, GradientTransformation(init=<function chain.<locals>.init_fn at 0x7f43f8875950>, update=<function chain.<locals>.update_fn at 0x7f43f8875a70>))], [*, CustomNode(<class 'flax.core.frozen_dict.FrozenDict'>[()], [{'params': {'Dense_0': {'bias': *, 'kernel': *}, 'Dense_1': {'bias': *, 'kernel': *}}}]), [CustomNode(namedtuple[<class 'optax._src.transform.ScaleByAdamState'>], [*, CustomNode(<class 'flax.core.frozen_dict.FrozenDict'>[()], [{'params': {'Dense_0': {'bias': *, 'kernel': *}, 'Dense_1': {'bias': *, 'kernel': *}}}]), CustomNode(<class 'flax.core.frozen_dict.FrozenDict'>[()], [{'params': {'Dense_0': {'bias': *, 'kernel': *}, 'Dense_1': {'bias': *, 'kernel': *}}}])]), CustomNode(namedtuple[<class 'flax.serialization.EmptyState'>], []), CustomNode(namedtuple[<class 'flax.serialization.EmptyState'>], [])]])) and PyTreeDef(CustomNode(<class 'flax.training.train_state.TrainState'>[(<bound method Module.apply of MLP(
    # attributes
    features = [4, 1]
)>, GradientTransformation(init=<function chain.<locals>.init_fn at 0x7f43f8875950>, update=<function chain.<locals>.update_fn at 0x7f43f8875a70>))], [*, CustomNode(<class 'flax.core.frozen_dict.FrozenDict'>[()], [{'params': {'Dense_0': {'bias': *, 'kernel': *}, 'Dense_1': {'bias': *, 'kernel': *}}}]), [CustomNode(namedtuple[<class 'flax.serialization.ScaleByAdamState'>], [*, CustomNode(<class 'flax.core.frozen_dict.FrozenDict'>[()], [{'params': {'Dense_0': {'bias': *, 'kernel': *}, 'Dense_1': {'bias': *, 'kernel': *}}}]), CustomNode(<class 'flax.core.frozen_dict.FrozenDict'>[()], [{'params': {'Dense_0': {'bias': *, 'kernel': *}, 'Dense_1': {'bias': *, 'kernel': *}}}])]), CustomNode(namedtuple[<class 'flax.serialization.EmptyState'>], []), CustomNode(namedtuple[<class 'flax.serialization.EmptyState'>], [])]])).

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------


The above exception was the direct cause of the following exception:

TypeError                                 Traceback (most recent call last)

<ipython-input-4-add6ba3b684b> in train_step(state, batch)
     23         return state
     24 
---> 25     new_state = jax.lax.cond(True, upd_step, no_upd_step, None)
     26     metrics = {"loss": loss}
     27 

TypeError: true_fun and false_fun output must have same type structure, got PyTreeDef(CustomNode(<class 'flax.training.train_state.TrainState'>[(<bound method Module.apply of MLP(
    # attributes
    features = [4, 1]
)>, GradientTransformation(init=<function chain.<locals>.init_fn at 0x7f43f8875950>, update=<function chain.<locals>.update_fn at 0x7f43f8875a70>))], [*, CustomNode(<class 'flax.core.frozen_dict.FrozenDict'>[()], [{'params': {'Dense_0': {'bias': *, 'kernel': *}, 'Dense_1': {'bias': *, 'kernel': *}}}]), [CustomNode(namedtuple[<class 'optax._src.transform.ScaleByAdamState'>], [*, CustomNode(<class 'flax.core.frozen_dict.FrozenDict'>[()], [{'params': {'Dense_0': {'bias': *, 'kernel': *}, 'Dense_1': {'bias': *, 'kernel': *}}}]), CustomNode(<class 'flax.core.frozen_dict.FrozenDict'>[()], [{'params': {'Dense_0': {'bias': *, 'kernel': *}, 'Dense_1': {'bias': *, 'kernel': *}}}])]), CustomNode(namedtuple[<class 'flax.serialization.EmptyState'>], []), CustomNode(namedtuple[<class 'flax.serialization.EmptyState'>], [])]])) and PyTreeDef(CustomNode(<class 'flax.training.train_state.TrainState'>[(<bound method Module.apply of MLP(
    # attributes
    features = [4, 1]
)>, GradientTransformation(init=<function chain.<locals>.init_fn at 0x7f43f8875950>, update=<function chain.<locals>.update_fn at 0x7f43f8875a70>))], [*, CustomNode(<class 'flax.core.frozen_dict.FrozenDict'>[()], [{'params': {'Dense_0': {'bias': *, 'kernel': *}, 'Dense_1': {'bias': *, 'kernel': *}}}]), [CustomNode(namedtuple[<class 'flax.serialization.ScaleByAdamState'>], [*, CustomNode(<class 'flax.core.frozen_dict.FrozenDict'>[()], [{'params': {'Dense_0': {'bias': *, 'kernel': *}, 'Dense_1': {'bias': *, 'kernel': *}}}]), CustomNode(<class 'flax.core.frozen_dict.FrozenDict'>[()], [{'params': {'Dense_0': {'bias': *, 'kernel': *}, 'Dense_1': {'bias': *, 'kernel': *}}}])]), CustomNode(namedtuple[<class 'flax.serialization.EmptyState'>], []), CustomNode(namedtuple[<class 'flax.serialization.EmptyState'>], [])]])).
```

### Steps to reproduce:
Linked [colab notebook](https://colab.research.google.com/drive/1J-mK1cWSunKCO9NgBA_Pb90PZcOTWgCi?usp=sharing) includes some workarounds. Below is code to reproduce:

```py
import jax
import jax.numpy as jnp
import flax
import flax.linen as nn
import optax

from typing import Sequence
from flax.training.train_state import TrainState
from flax.training.checkpoints import save_checkpoint, restore_checkpoint


rng = jax.random.PRNGKey(842)
rng, data_rng = jax.random.split(rng)
x = jnp.array([[x, x] for x in range(64)], dtype=jnp.float32)
y = jnp.sum(2*x + 1, axis=-1, keepdims=True)
x = x + jax.random.normal(data_rng, x.shape)
def data_gen():
    yield x, y
    
class MLP(nn.Module):
    features: Sequence[int]

    @nn.compact
    def __call__(self, x):
        for feat in self.features[:-1]:
            x = nn.relu(nn.Dense(feat)(x))
        x = nn.Dense(self.features[-1])(x)
        return x

model = MLP([4, 1])
params = model.init(jax.random.PRNGKey(0), x)

optimizer = optax.adamw(0.01)
optimizer = optax.MultiSteps(optimizer, 4)
state = TrainState.create(apply_fn=model.apply, params=params, tx=optimizer)


def compute_loss(params, batch):
    preds = state.apply_fn(params, batch[0])
    targs = batch[1]
    return jnp.mean((preds - targs)**2)

grad_fn = jax.value_and_grad(compute_loss)

def train_step(state, batch):
    def compute_loss(params):
        preds = state.apply_fn(params, batch[0])
        targs = batch[1]
        return jnp.mean((preds - targs)**2)
    grad_fn = jax.value_and_grad(compute_loss)
    
    loss, grad = grad_fn(state.params)

    new_state = state.apply_gradients(grads=grad)

    metrics = {"loss": loss}

    return new_state, metrics
train_step = jax.jit(train_step)

# train model, save checkpoint
for i in range(8):
    batch = next(data_gen())
    state, metrics = train_step(state, batch)
    print(metrics["loss"])
save_checkpoint('./_tmp/', state, 8, overwrite=True)

# restore checkopint, resume training - fails
state = restore_checkpoint('./_tmp/', state)
for i in range(8):
    batch = next(data_gen())
    state, metrics = train_step(state, batch)
    print(metrics["loss"])
```



Discussion:
This is definitely a bug on our side. We are essentially recreating named tuples but I'm surprised that didn't result in trouble before. I will fix this.

Resolving patch:
diff --git a/flax/serialization.py b/flax/serialization.py
--- a/flax/serialization.py
+++ b/flax/serialization.py
@@ -22,6 +22,7 @@
 import jax
 import msgpack
 import numpy as np
+from numpy.lib.arraysetops import isin
 
 
 _STATE_DICT_REGISTRY = {}
@@ -125,27 +126,24 @@ def _restore_dict(xs, states):
 
 
 def _namedtuple_state_dict(nt):
-  return {'name': nt.__class__.__name__,
-          'fields': {str(i): to_state_dict(x)
-                     for i, x in enumerate(nt._fields)},
-          'values': {str(i): to_state_dict(x)
-                     for i, x in enumerate(nt)}
-         }
+  return {key: to_state_dict(getattr(nt, key)) for key in nt._fields}
 
 
 def _restore_namedtuple(xs, state_dict):
   """Rebuild namedtuple from serialized dict."""
-  if len(state_dict['values']) != len(xs):
-    raise ValueError('The size of the list and the state dict do not match,'
-                     f' got {len(xs)} and {len(state_dict["values"])}.')
-  fields = [state_dict['fields'][str(i)] for i in range(len(xs))]
-  namedtuple_class = collections.namedtuple(
-      state_dict['name'], fields)
-  ys = []
-  for i in range(len(state_dict['values'])):
-    y = from_state_dict(xs[i], state_dict['values'][str(i)])
-    ys.append(y)
-  return namedtuple_class(*ys)
+  if set(state_dict.keys()) == {'name', 'fields', 'values'}:
+    # TODO(jheek): remove backward compatible named tuple restoration early 2022
+    state_dict = {state_dict['fields'][str(i)]: state_dict['values'][str(i)]
+                  for i in range(len(state_dict['fields']))}
+
+  sd_keys = set(state_dict.keys())
+  nt_keys = set(xs._fields)
+
+  if sd_keys != nt_keys:
+    raise ValueError('The field names of the state dict and the named tuple do not match,'
+                     f' got {sd_keys} and {nt_keys}.')
+  fields = {k: from_state_dict(getattr(xs, k), v) for k, v in state_dict.items()}
+  return type(xs)(**fields)
 
 
 register_serialization_state(dict, _dict_state_dict, _restore_dict)


Test patch:
diff --git a/tests/serialization_test.py b/tests/serialization_test.py
--- a/tests/serialization_test.py
+++ b/tests/serialization_test.py
@@ -212,6 +212,20 @@ def test_namedtuple_serialization(self):
     x1_serialized = serialization.to_bytes(x1)
     x2 = foo_class(a=0, b=0, c=0)
     restored_x1 = serialization.from_bytes(x2, x1_serialized)
+    self.assertEqual(type(x1), type(restored_x1))
+    self.assertEqual(x1, restored_x1)
+  
+  def test_namedtuple_restore_legacy(self):
+    foo_class = collections.namedtuple('Foo', 'a b c')
+    x1 = foo_class(a=1, b=2, c=3)
+    legacy_encoding = {
+      'name': 'Foo',
+      'fields': {'0': 'a', '1': 'b', '2': 'c'},
+      'values': {'0': 1, '1': 2, '2': 3},
+    }
+    x2 = foo_class(a=0, b=0, c=0)
+    restored_x1 = serialization.from_state_dict(x2, legacy_encoding)
+    self.assertEqual(type(x1), type(restored_x1))
     self.assertEqual(x1, restored_x1)
 
   def test_model_serialization_to_bytes(self):