PR #5433 [SPMD] Make IR sharding custom sharding op

[SPMD] Make IR sharding custom sharding op #5433

alanwaketan merged 20 commits into master from alanwaketan/custom_sharding

alanwaketan requested a review from

yeounoh 2 years ago

alanwaketan requested a review from

jonb377 2 years ago

alanwaketan requested a review from

JackCaoG 2 years ago

alanwaketan assigned

alanwaketan 2 years ago

JackCaoG added distributed

JackCaoG commented on 2023-08-10

alanwaketan changed the base branch from master to alanwaketan/debug 2 years ago

alanwaketan requested a review from

will-cromar 2 years ago

alanwaketan requested a review from

mateuszlewko 2 years ago

alanwaketan requested a review from

stgpetrovic 2 years ago

alanwaketan changed the base branch from alanwaketan/debug to master 2 years ago

alanwaketan removed review request from

stgpetrovic 2 years ago

alanwaketan removed review request from

mateuszlewko 2 years ago

alanwaketan removed review request from

will-cromar 2 years ago

JackCaoG commented on 2023-08-10

alanwaketan force pushed from ffcc94c7 to 501b96f9 2 years ago

alanwaketan changed the base branch from master to alanwaketan/debug 2 years ago

alanwaketan changed the base branch from alanwaketan/debug to master 2 years ago

Fix linters

1036b93f

Fix linters

5459ca17

Add lowering

43047804

add ir node

1ba36186

add tensor methods

ac26dafa

connect the pybind

9622c475

Modify test and small fix

1415d565

Fix the custom call hlo

c4ff7d92

Fix test

2d468d52

Don't use XLAShardedTensor

e45ec012

Make it in-place op

da528248

Fix tests

63a5734d

Fix the test

3c775212

Do custom call in _propagate_xla_data

605b46e8

Fix linters

6cbf580f

Fix linters

7c57ee1c

Revert vscode init

97201827

minor fix

2cb25577

Make custom_sharding_

10b07b70

Fix linters

743978fe

alanwaketan force pushed from a1141042 to 743978fe 2 years ago

JackCaoG commented on 2023-08-23

JackCaoG approved these changes on 2023-08-24

alanwaketan merged dea9b296 into master 2 years ago

yeounoh commented on 2023-08-28

Reviewers

JackCaoG

yeounoh

jonb377

Assignees

alanwaketan

Labels

distributed

Milestone

No milestone

xla [SPMD] Make IR sharding custom sharding op #5433 Merged

[SPMD] Make IR sharding custom sharding op #5433

xla
[SPMD] Make IR sharding custom sharding op
#5433

Merged