Add `torch.count_nonzero` (#39992)
Summary:
Reference https://github.com/pytorch/pytorch/issues/38349
TODO:
* [x] Add tests
* [x] Add docs (pending add to docs.rst)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/39992
Reviewed By: ezyang
Differential Revision: D22236738
Pulled By: mruberry
fbshipit-source-id: 8520068b086b5ffc4de9e4939e746ff889293987