flax

Call user-defined variable transforms before determining axis size in nn.vmap.

#4026

Merged