Experiment with using out_axes=-1 inside nmap's vmap#99
Draft
shoyer wants to merge 1 commit intogoogle-deepmind:mainfrom
Draft
Experiment with using out_axes=-1 inside nmap's vmap#99shoyer wants to merge 1 commit intogoogle-deepmind:mainfrom
shoyer wants to merge 1 commit intogoogle-deepmind:mainfrom
