flax

Avoid passing non-boolean mask to `where` argument of `jax.numpy` reductions. Non-boolean mask inputs have been deprecated for several releases, and will result in an error starting in JAX v0.8.0.

#4923

Merged