Experiment with using out_axes=-1 inside nmap's vmap#99

Draft

shoyer wants to merge 1 commit intogoogle-deepmind:mainfrom shoyer:vmap-end

+8-5

Commits on Nov 25, 2024

Experiment with using out_axes=-1 inside nmap's vmap
shoyer
committed