flax

Fix jax.tree_util.register_dataclass in older JAX versions.

#3885

Merged