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

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

samanklesaria wants to merge 2 commits into google:main from samanklesaria:handle-captures-jit

samanklesaria force pushed from ff0a4983 to b1aeb436 6 days ago

samanklesaria force pushed from b1aeb436 to cd637914 6 days ago

samanklesaria force pushed from cd637914 to bae7839b 6 days ago

samanklesaria force pushed from bae7839b to 62a03606 6 days ago

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

samanklesaria requested a review from

cgarciae 6 days ago

samanklesaria commented on 2026-06-16

optimize jit_partial

66ecd95b

IvyZX commented on 2026-06-16

samanklesaria force pushed from 956b8075 to 45c8ac05 1 day ago

Handle captures in jax transforms

43601424

samanklesaria force pushed from 45c8ac05 to 43601424 1 day ago

Reviewers

IvyZX

cgarciae

Assignees

No one assigned

Labels

None yet

Milestone

No milestone