flax
Fix NNX jit static args with in_shardings issue #4989
#4996

Merged

Fix NNX jit static args with in_shardings issue #4989 #4996

copybara-service merged 4 commits into google:main from mohsinm-dev:fix-nnx-jit-static-sharding

vfdev-5 commented on 2025-10-03

samanklesaria commented on 2025-10-05

mohsinm-dev requested a review from

samanklesaria 74 days ago

samanklesaria commented on 2025-10-08

vfdev-5 commented on 2025-10-23

vfdev-5 commented on 2025-10-27

mohsinm-dev force pushed from 1dd4306c to 7a86859f 54 days ago

vfdev-5 commented on 2025-10-28

vfdev-5 requested a review from

cgarciae 50 days ago

cgarciae approved these changes on 2025-11-03

cgarciae added pull ready

Fix NNX jit static args with in_shardings issue #4989

963f9ba9

Remove redundant static args sharding tests

6840519c

Fix trailing whitespace in transforms_test.py

b9ac1bce

vfdev-5 force pushed from d4cdec45 to 38285f74 46 days ago

Removed private jax._src.api_util import and copied them into flax

dcf120b6

vfdev-5 force pushed from 38285f74 to dcf120b6 46 days ago

cgarciae approved these changes on 2025-11-06

copybara-service merged d3754f63 into main 44 days ago

Reviewers

cgarciae

vfdev-5

samanklesaria

Assignees

No one assigned

Labels

pull ready

Milestone

No milestone