PR #5231 Add 'opt_sharding' option for Variables

Add 'opt_sharding' option for Variables #5231

samanklesaria wants to merge 5 commits into google:main from samanklesaria:opt_sharding

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

samanklesaria force pushed from f11d1cf5 to 5e4ca369 26 days ago

samanklesaria marked this pull request as ready for review 26 days ago

samanklesaria requested a review from

cgarciae 25 days ago

Add 'opt_sharding' option for Variables

e08e3d1c

Update flax/nnx/training/optimizer.py

82a3c1b7

samanklesaria force pushed from 20a4ae43 to 82a3c1b7 3 days ago

IvyZX approved these changes on 2026-03-04

cgarciae commented on 2026-03-04

Update out_sharding name

8520c75b

Fix linter errors

0966c0e7

Update name to optimizer_sharding

61e0d32d

Reviewers

IvyZX

cgarciae

gemini-code-assist

Assignees

No one assigned

Labels

None yet

Milestone

No milestone