flax

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