PR #5231 Add 'opt_sharding' option for Variables

Add 'opt_sharding' option for Variables #5231

copybara-service merged 5 commits into google:main from samanklesaria:opt_sharding

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

samanklesaria force pushed from f11d1cf5 to 5e4ca369 204 days ago

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

samanklesaria requested a review from

cgarciae 203 days ago

Add 'opt_sharding' option for Variables

e08e3d1c

Update flax/nnx/training/optimizer.py

82a3c1b7

samanklesaria force pushed from 20a4ae43 to 82a3c1b7 181 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

cgarciae approved these changes on 2026-03-24

cgarciae added pull ready

copybara-service merged bb50f212 into main 161 days ago

Reviewers

cgarciae

IvyZX

gemini-code-assist

Assignees

No one assigned

Labels

pull ready

Milestone

No milestone