flax

7cca4073 - ignore optax deprecation warning