Bo Wang b6df043f1f Add torch.nn.init.uniform_ operator to ShardedTensor. (#63997) ... Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/63997 Use torch_function to extend torch.nn.init.uniform_ The Init is done in SPMD fashion. Note that ideally we want to aggregate sharded tensors into a global tensor, init it and reshard. It's fine to run it SPMD since uniform is I.I.D indepenent and identifically distributed. Also enable unit test for test_linear.py for OSS test Test Plan: a) Unit Test (pytorch) ... $ python test/distributed/_sharded_tensor/ops/test_init.py TestShardedTensorNNInit --v (pytorch) ... $ python test/distributed/_sharded_tensor/ops/test_linear.py --v (before runs this command is no-op) or b) Manual run: Instruction here: https://docs.google.com/document/d/1_m1Hdo5w51-hhPlZ_F8Y6PIWrN7UgJZqiSpARYvhsaE/edit# Imported from OSS Reviewed By: pritamdamania87, anjali411 Differential Revision: D30563017 fbshipit-source-id: d1859f7682235bcb44515efc69ca92bc5e34fce1		2021-10-21 00:17:13 -07:00
..
ops	Add torch.nn.init.uniform_ operator to ShardedTensor. (#63997)	2021-10-21 00:17:13 -07:00
__init__.py	Basic implementation of ShardedLinear using ShardedTensor. (#64128)	2021-09-20 18:31:11 -07:00
api.py	Add torch.nn.init.uniform_ operator to ShardedTensor. (#63997)	2021-10-21 00:17:13 -07:00