xla
Add API to assemble CPU shards to a sharded tensor
#5681

Merged

Add API to assemble CPU shards to a sharded tensor #5681

jonb377 merged 6 commits into master from jonbolin-assemble-shards

Add API to assemble CPU shards to a sharded tensor

dd6a770f

Handle replicated sharding

ca0ee017

Move validations into get_op_sharding

23cfcc44

Improve tests and error handling

e4e7ffd5

Don't WrapXlaData

31793b28

jonb377 requested a review from

alanwaketan 2 years ago

alanwaketan approved these changes on 2023-10-05

Fix test for v3

371f99c6

jonb377 merged d9a90499 into master 2 years ago

jonb377 deleted the jonbolin-assemble-shards branch 2 years ago

Reviewers

alanwaketan

Assignees

No one assigned

Labels

None yet

Milestone

No milestone