Handle captures in jax transforms#5485

samanklesaria wants to merge 1 commit into

google:maingoogle/flax:mainfrom

samanklesaria:handle_capturessamanklesaria/flax:handle_capturesCopy head branch name to clipboard