Handle capture variables in nnx.jit. #5490

samanklesaria wants to merge 1 commit into

google:maingoogle/flax:mainfrom

samanklesaria:handle-captures-jitsamanklesaria/flax:handle-captures-jitCopy head branch name to clipboard