flax

Save types information to allow for fully self-describing checkpoints. Add type-specific metadata to `LazyValue` so that more information can be gleaned without fully restoring the tree.

#3136

Open