General fixes for ShardedTensor op framework.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/77191
1) Add more basic validation to all ops.
2) Ensure `register_on_local_shards` uses appropriate sharding_spec.
Differential Revision: [D36292159](https://our.internmc.facebook.com/intern/diff/D36292159/)
Approved by: https://github.com/wanchaol, https://github.com/fduwjj
