Handle captures in nnx.vmap#5527

samanklesaria wants to merge 4 commits into

google:maingoogle/flax:mainfrom

samanklesaria:handle_captures_vmapsamanklesaria/flax:handle_captures_vmapCopy head branch name to clipboard