flax
Added out_sharding and qkv_sharding to MHA call method
#5434

Merged

Added out_sharding and qkv_sharding to MHA call method #5434

copybara-service merged 2 commits into main from mha-out_sharding

Added out_sharding and qkv_sharding to MHA call method

42edb5be

samanklesaria approved these changes on 2026-04-23

vfdev-5 marked this pull request as ready for review 63 days ago

cgarciae approved these changes on 2026-04-23

Merge branch 'main' into mha-out_sharding

502012b8

cgarciae added pull ready

copybara-service merged 2fa7e66e into main 45 days ago

copybara-service deleted the mha-out_sharding branch 45 days ago

Reviewers

cgarciae

samanklesaria

Assignees

No one assigned

Labels

pull ready

Milestone

No milestone