add type annotations to comm.py (#46736)
Summary:
Fixes https://github.com/pytorch/pytorch/issues/46735
Pull Request resolved: https://github.com/pytorch/pytorch/pull/46736
Reviewed By: albanD
Differential Revision: D24565554
Pulled By: mrshenli
fbshipit-source-id: 4e40e4232ebf256af228f9c742ea4d28c626c616
