flax
Fix for change to `jax.Array.__contains__`
#5404
Merged

Loading