flax
feat(nnx): add out_sharding to recurrent cells
#5255

Open

feat(nnx): add out_sharding to recurrent cells #5255

aarushisingh04 wants to merge 4 commits into google:main from aarushisingh04:pr3-recurrent-out-sharding

gemini-code-assist commented on 2026-02-15

feat(nnx): Add out_sharding to recurrent cell __call__ methods

fcca8e94

aarushisingh04 force pushed from 6e465017 to fcca8e94 129 days ago

gemini-code-assist commented on 2026-02-15

Update flax/nnx/nn/recurrent.py

331ff8e5

samanklesaria approved these changes on 2026-04-21

IvyZX approved these changes on 2026-04-23

IvyZX added pull ready

IvyZX requested changes on 2026-04-23

IvyZX removed pull ready

test(nnx): verify out_sharding applies correct sharding to recurrent …

261ddfb4

IvyZX requested changes on 2026-04-28

fix: use meaningful sharding spec and device-count batch size in out_…

0a18acc4

Reviewers

samanklesaria

IvyZX

gemini-code-assist

Assignees

No one assigned

Labels

None yet

Milestone

No milestone