flax
Add out_sharding argument to call methods for layers with jax calls that support it
#5102

Merged

Add out_sharding argument to call methods for layers with jax calls that support it #5102

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

samanklesaria force pushed from 67bafc18 to 11950870 110 days ago

samanklesaria force pushed from 11950870 to 96f0a646 110 days ago

Add out_sharding argument to call methods for standard layers

1f61d6b7

samanklesaria force pushed from 96f0a646 to 1f61d6b7 110 days ago

samanklesaria requested a review from

cgarciae 110 days ago

samanklesaria changed the title ~~Add out_sharding argument to call methods for standard layers~~ Add out_sharding argument to call methods for layers with jax calls that support it 110 days ago

cgarciae commented on 2025-11-25

cgarciae approved these changes on 2025-11-25

cgarciae added pull ready

chapman20j commented on 2025-11-25

Fix nits

66544200

chapman20j approved these changes on 2025-12-11

copybara-service merged ef78d658 into main 88 days ago

Reviewers

chapman20j

cgarciae

Assignees

No one assigned

Labels

pull ready

Milestone

No milestone

flax Add out_sharding argument to call methods for layers with jax calls that support it #5102 Merged

Add out_sharding argument to call methods for layers with jax calls that support it #5102

flax
Add out_sharding argument to call methods for layers with jax calls that support it
#5102

Merged