Compatibility with newest Jax and Flax versions#75

Merged

markusschmitt merged 4 commits intomasterfrom jax_flax_version_compatibility

Sep 11, 2024

+45-47

Commits on Sep 11, 2024

Fixed issues with new Jax/Flax versions.
Markus Schmitt
committed
Version number.
Markus Schmitt
committed
Fixes for backward compatibility.
Markus Schmitt
committed
Update.
Markus Schmitt
committed