flax

90d656a8 - [nnx] improve vmap axis size detection