flax
make flax work with upcoming JAX change to tree_map (being more careful about
#3976
Merged
Merged