flax
Fix ValueError when `nnx.jit` is used with `nnx.custom_vjp`
#5045

Merged

Fix ValueError when `nnx.jit` is used with `nnx.custom_vjp` #5045

copybara-service merged 1 commit into google:main from samanklesaria:issues/5044

samanklesaria changed the title ~~Prevent child ctxtag from overwriting parent~~ Fix ValueError when `nnx.jit` is used with `nnx.custom_vjp` 68 days ago

samanklesaria requested a review from

vfdev-5 68 days ago

vfdev-5 commented on 2025-10-23

samanklesaria force pushed from 2bc02fe1 to 79b0dd4a 67 days ago

samanklesaria requested a review from

vfdev-5 67 days ago

vfdev-5 commented on 2025-10-23

vfdev-5 approved these changes on 2025-10-23

samanklesaria requested a review from

cgarciae 67 days ago

cgarciae approved these changes on 2025-10-23

cgarciae added pull ready

Prevent child ctxtag from overwriting parent

fb391c11

samanklesaria force pushed from c23ea853 to fb391c11 66 days ago

copybara-service merged dc7e7088 into main 66 days ago

Reviewers

cgarciae

vfdev-5

Assignees

No one assigned

Labels

pull ready

Milestone

No milestone

flax Fix ValueError when `nnx.jit` is used with `nnx.custom_vjp` #5045 Merged

Fix ValueError when `nnx.jit` is used with `nnx.custom_vjp` #5045

flax
Fix ValueError when `nnx.jit` is used with `nnx.custom_vjp`
#5045

Merged