flax
Generalize out_sharding to work with NamedSharding and Format
#5246

Merged

Generalize out_sharding to work with NamedSharding and Format #5246

copybara-service merged 2 commits into google:main from samanklesaria:generalize_out_sharding

samanklesaria requested a review from

cgarciae 44 days ago

samanklesaria commented on 2026-02-13

cgarciae commented on 2026-03-04

samanklesaria requested a review from

cgarciae 23 days ago

Generalize out_sharding to work with NamedSharding and Format

0399b9d3

Use with_sharding_constraint for Format application

1c2cec46

samanklesaria force pushed from 52ce6753 to 1c2cec46 10 days ago

cgarciae approved these changes on 2026-03-19

cgarciae added pull ready

copybara-service merged 3bcc4a91 into main 9 days ago

Reviewers

cgarciae

Assignees

No one assigned

Labels

pull ready

Milestone

No milestone