pytorch
3ee1f81d - Extend _sharded_tensor constructor to support other ops like torch.ones (#63378)

Commit
4 years ago
Extend _sharded_tensor constructor to support other ops like torch.ones (#63378) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/63378 a) Introduce InitCommonParams to wrap tensor creation params b) Factor local tensor initiation into common_params so that tensor value is not hard specified in ShardedTensor constructor c) Add _sharded_tensor.ones(...) to exemplify - Note memory_format arg is not provided to be consistent as torch.ones d) Follow up: more ops like torch.full, torch.zero, torch.rand, Test: $ python test/distributed/_sharded_tensor/test_sharded_tensor.py TestCreateTensorFromParams --v $ python test/distributed/_sharded_tensor/test_sharded_tensor.py TestShardedTensorChunked.test_create_sharded_tensor_with_ones --v $ python test/distributed/_sharded_tensor/test_sharded_tensor.py TestShardedTensorEnumerable.test_create_sharded_tensor_with_ones --v Test Plan: Imported from OSS Reviewed By: pritamdamania87, wanchaol Differential Revision: D30359245 Pulled By: bowangbj fbshipit-source-id: 85768fcb36e9d9d40213036884b1266930a91701
Author
Parents
Loading