flax

bfca2cc2 - Fix serialization restore.

Fix serialization restore. Serialization was creating a new named tuple class based on the fields in the state dict. This results in a type error when comparing to the original or when comparing the structure of PyTrees.