pytorch

3596e13d - Add torch.nn.init.normal_ and torch.nn.init.kaiming_uniform_ ops to ShardedTensor (#67057)

Add torch.nn.init.normal_ and torch.nn.init.kaiming_uniform_ ops to ShardedTensor (#67057) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/67057 Extend ShardedTensor with torch.nn.init.[normal_, and kaiming_uniform_] ops Follow up from https://github.com/pytorch/pytorch/pull/63997 Test Plan: a) Unit Test (pytorch) ... $ python test/distributed/_sharded_tensor/ops/test_init.py TestShardedTensorNNInit --v or b) Manual run: Instruction here: https://docs.google.com/document/d/1_m1Hdo5w51-hhPlZ_F8Y6PIWrN7UgJZqiSpARYvhsaE/edit# s/uniform_/normal_ or kaiming_uniform_ Imported from OSS Reviewed By: pritamdamania87 Differential Revision: D31845654 fbshipit-source-id: e7aedc0972539da59f7b84bbbf617caf6b206d52