Merge global_device_attributes

Author: will-cromar

test/pjrt/test_runtime_tpu.py

@@ -207,10 +207,10 @@ def test_device_attributes(self):
 
   @staticmethod
   def _global_device_attributes():
-    return pjrt.global_device_attributes()
+    return xr.global_device_attributes()
 
   def test_global_device_attributes(self):
-    results = pjrt._run_multiprocess(self._global_device_attributes)
+    results = pjrt.run_multiprocess(self._global_device_attributes)
     for result in results.values():
       for device in result:
         self.assertCountEqual(['coords', 'core_on_chip'], list(device.keys()))

torch_xla/experimental/pjrt.py

@@ -9,6 +9,7 @@
     runtime.addressable_device_count,
     runtime.device_attributes,
     runtime.device_type,
+    runtime.global_device_attributes,
     runtime.global_device_count,
     runtime.global_ordinal,
     runtime.local_device_count,