gemma icon indicating copy to clipboard operation
gemma copied to clipboard

Remove pinned version of `flax`

Open DwarKapex opened this issue 2 years ago • 1 comments

Hi everyone,

Is it necessary for the project to pin the version of flax? Can you use the latest flax or to use version greater or less than 0.7.5?

Thank you in advance for the reply and fix if possible =)

DwarKapex avatar Feb 26 '24 19:02 DwarKapex

https://github.com/google/flax/blob/v0.7.5/flax/configurations.py#L42 is not compatible with recent JAX: https://github.com/google/jax/commit/57e59eb6c309acfdfb49a4ce76c996fc5c2a016c.

olupton avatar Mar 01 '24 09:03 olupton