allow captures in nnx.grad#5484

copybara-service[bot] wants to merge 1 commit into

maingoogle/flax:mainfrom

test_903485800google/flax:test_903485800Copy head branch name to clipboard