flax
fb771186 - Add vmap_with_axes, a logical axis version of nn.vmap

Commit

3 years ago

Add vmap_with_axes, a logical axis version of nn.vmap PiperOrigin-RevId: 424963874

Author

James Lee-Thorp

Committer

a-googler

Parents