flax
Add out_sharding argument to Embed layer call
#5205

Merged

Add out_sharding argument to Embed layer call #5205

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

Add out_sharding argument to Embed layer call

ca19b9bc

samanklesaria requested a review from

cgarciae 46 days ago

gemini-code-assist commented on 2026-01-27

Apply suggestions from code review

bec643f7

Revert test to use Exception rather than ValueError

ce2ae486

IvyZX commented on 2026-02-02

IvyZX approved these changes on 2026-02-04

IvyZX added pull ready

copybara-service merged 2e1f5adc into main 38 days ago

Reviewers

IvyZX

gemini-code-assist

cgarciae

Assignees

No one assigned

Labels

pull ready

Milestone

No milestone