xla
ffa5b340 - [SPMD][PoC] InputHandler for input sharding (#3963)

Commit
3 years ago
[SPMD][PoC] InputHandler for input sharding (#3963) * [SPMD] Add at::Tensor sharding util & sharding tests. * [SPMD] Add InputHandler for input sharding. * [SPMD] Add PjRtShardedData & sharded data helper functions * [SPMD] Add CreateTensorsData to shard & transfer tensor data * [SPMD] Use shared_ptr for output_sharding_ in Node * [SPMD] attach sharding annotatino to IR node only. * Disable test_xla_sharding until we land all POC PRs * [SPMD] Refactoring
Author
Parents
Loading