flax
Add out_sharding arguments to linear layers where supported
#5156

Merged

Add out_sharding arguments to linear layers where supported #5156

copybara-service merged 2 commits into google:main from jackopenn:conv-out-shardings

samanklesaria requested a review from

samanklesaria 8 days ago

samanklesaria approved these changes on 2025-12-29

jackopenn requested a review from

samanklesaria 8 days ago

samanklesaria approved these changes on 2025-12-29

Add out_sharding arguments to linear layers where supported

23a1c1b1

fix: update expected shape in conv sharding test

53fbca51

samanklesaria force pushed from 6c8b32e6 to 53fbca51 1 day ago

IvyZX approved these changes on 2026-01-05

IvyZX added pull ready

copybara-service merged c5398d67 into main 1 day ago

Reviewers

IvyZX

samanklesaria

Assignees

No one assigned

Labels

pull ready

Milestone

No milestone