xla
[SPMD] Support reduce-scatter in manual sharding
#7231

Merged

[SPMD] Support reduce-scatter in manual sharding #7231

alanwaketan merged 9 commits into master from alanwaketan/spmd_asr

initial commit

12b5a7ed

Add a test case

ad5b7a66

fix the lowering

38b28e70

add canonical index

6f4f1709

Fix test

bccdacd9

Add more test

aa68e604

Fix linters

a212189c

alanwaketan added tpuci

alanwaketan requested a review from

jonb377 1 year ago

alanwaketan requested a review from

JackCaoG 1 year ago

alanwaketan assigned

alanwaketan 1 year ago

Fix linters

fec38604

JackCaoG commented on 2024-06-10

Skip non TPU devices

48eef410

JackCaoG approved these changes on 2024-06-11

alanwaketan merged 70d2e9ea into master 1 year ago

alanwaketan added backport_2.4

Reviewers

JackCaoG

jonb377

Assignees

alanwaketan

Labels

backport_2.4

Milestone

No milestone