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

References

Author

epiqueras

Committer

a-googler

Parents