Handle captures in jax transforms with hijax - #5531

Handle captures in jax transforms with hijax#5531

samanklesaria wants to merge 1 commit into

google:maingoogle/flax:mainfrom

samanklesaria:handle-captures-jit-hijaxsamanklesaria/flax:handle-captures-jit-hijaxCopy head branch name to clipboard