[PJRT] Update `torch.distributed` unit test to use PJRT (#4743)
* Start refactoring torch.distributed test to support PJRT
* Fix imports
* Use unit test assertions
* pjrt init method
* Add back class method
* Formatting
* Fix Send and Recv tests
* Remove unused import
* Use dist.send/recv
