flax

Allow restoring raw state_dict using checkpoints.restore_checkpoint().

#563

Merged