PR #5490 Handle capture variables in `nnx.jit`.

Handle capture variables in `nnx.jit`. #5490

samanklesaria wants to merge 1 commit into google:main from samanklesaria:handle-captures-jit

samanklesaria force pushed from ff0a4983 to b1aeb436 51 days ago

samanklesaria force pushed from b1aeb436 to cd637914 51 days ago

samanklesaria force pushed from cd637914 to bae7839b 51 days ago

samanklesaria force pushed from bae7839b to 62a03606 51 days ago

samanklesaria force pushed from 62a03606 to 956b8075 51 days ago

samanklesaria requested a review from

cgarciae 51 days ago

samanklesaria commented on 2026-06-16

IvyZX commented on 2026-06-16

samanklesaria force pushed from 956b8075 to 45c8ac05 46 days ago

samanklesaria force pushed from 45c8ac05 to 43601424 46 days ago

samanklesaria force pushed from 43601424 to 5ec20cc3 44 days ago

samanklesaria force pushed from 5ec20cc3 to 97e979de 44 days ago

cgarciae commented on 2026-06-18

samanklesaria force pushed from 97e979de to 33e956e0 43 days ago

samanklesaria requested a review from

cgarciae 43 days ago

samanklesaria force pushed from 33e956e0 to df8075ac 43 days ago

samanklesaria force pushed from df8075ac to 3e65391b 40 days ago

samanklesaria force pushed from 3e65391b to e4cd1a8d 40 days ago

samanklesaria force pushed from e4cd1a8d to 66dd25a8 40 days ago

samanklesaria force pushed from 66dd25a8 to e3d56755 37 days ago

samanklesaria force pushed from e3d56755 to 5bbd1b5c 37 days ago

samanklesaria force pushed from 5bbd1b5c to c84bbd0e 32 days ago

samanklesaria force pushed from c84bbd0e to d3ced874 32 days ago

samanklesaria force pushed from d3ced874 to f116d3f4 30 days ago

samanklesaria force pushed from f116d3f4 to ee772299 30 days ago

samanklesaria force pushed from ee772299 to d8aa50da 30 days ago

Handle captures in jax transforms

174b20b5

samanklesaria force pushed from d8aa50da to 174b20b5 26 days ago

Reviewers

cgarciae

IvyZX

Assignees

No one assigned

Labels

None yet

Milestone

No milestone