flax

update doctest requirements to use recentish jax minver

#3160

Merged