PR #5079 Add nnx.use_eager_sharding context manager

Add nnx.use_eager_sharding context manager #5079

copybara-service merged 3 commits into google:main from samanklesaria:eager_sharding_context

samanklesaria force pushed from 6579ea64 to 319fd55f 44 days ago

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

cgarciae commented on 2025-11-07

samanklesaria force pushed from 319fd55f to ed017702 44 days ago

cgarciae approved these changes on 2025-11-07

samanklesaria force pushed from ed017702 to 27e20a1e 41 days ago

samanklesaria force pushed from 27e20a1e to 56d65efc 41 days ago

samanklesaria force pushed from 56d65efc to ea81c813 41 days ago

samanklesaria force pushed from ea81c813 to 0f0b57de 41 days ago

samanklesaria force pushed from 0f0b57de to 45087575 41 days ago

samanklesaria force pushed from 45087575 to dda9e5a6 41 days ago

Add nnx.use_eager_sharding context manager

f600dcb5

samanklesaria force pushed from dda9e5a6 to f600dcb5 40 days ago

Merge branch 'main' into eager_sharding_context

0f5cdf14

Update gspmd guide to mention Variable init with eager_sharding

719f6879

cgarciae added pull ready

copybara-service merged 7710c304 into main 39 days ago

Reviewers

cgarciae

Assignees

No one assigned

Labels

pull ready

Milestone

No milestone