PR #5522 Handle captures grad

Handle captures grad #5522

samanklesaria wants to merge 3 commits into google:main from samanklesaria:handle-captures-grad

samanklesaria marked this pull request as ready for review 28 days ago

samanklesaria requested a review from

cgarciae 28 days ago

Handle captures in jax transforms

174b20b5

Handle captures in scan transformations

4aaf86c7

Handle captures in nnx.grad transforms

f054bbc8

samanklesaria force pushed from a8194315 to f054bbc8 24 days ago

Reviewers

cgarciae

Assignees

No one assigned

Labels

None yet

Milestone

No milestone