flax
ea3bcab6 - Fix jax.tree_util.register_dataclass in older JAX versions.

Commit
1 year ago
Fix jax.tree_util.register_dataclass in older JAX versions. PiperOrigin-RevId: 628149376
Author
Committer
Parents
Loading