Handle captures in nnx.scan#5515

samanklesaria wants to merge 2 commits into

google:maingoogle/flax:mainfrom

samanklesaria:handle-captures-scansamanklesaria/flax:handle-captures-scanCopy head branch name to clipboard