Deprecate devkind field from xla_model.xla_device (#9284)

Author: ghpvnist

benchmarks/check_xla_device.py

@@ -17,7 +17,7 @@ def use_torch_xla2():
 
 if not use_torch_xla2():
   import torch_xla.core.xla_model as xm
-  devlist = xm.get_xla_supported_devices(devkind=devkind)
+  devlist = xm.get_xla_supported_devices()
 else:
   # torch_xla2 needs jax to detect the device
   os.environ["JAX_PLATFORMS"] = devkind.lower(

scripts/bench_tensor_io.py

@@ -10,7 +10,7 @@
 
 
 def run_benchmark(args, pos_args):
-  devices = xm.get_xla_supported_devices(max_devices=args.max_devices)
+  devices = xm.get_xla_supported_devices(args.max_devices)
   shape = [int(x) for x in args.shape.split(',')]
 
   send_list = []

test/ds/test_dynamic_shape_models.py

@@ -44,7 +44,7 @@ def forward(self, x):
 
 
 @unittest.skipIf(
-    not xm.get_xla_supported_devices("TPU"),
+    xm.xla_device_hw(torch_xla.device()) != 'TPU',
     f"The tests fail on CPU. See https://github.com/pytorch/xla/issues/4298 for more detail."
 )
 class TestDynamicShapeModels(unittest.TestCase):

test/pjrt/test_dynamic_plugin_tpu.py

@@ -20,7 +20,7 @@ def setUpClass(cls):
   @staticmethod
   def _assert_tpus_exist(index=0):
     del index
-    assert len(xm.get_xla_supported_devices('TPU')) > 0
+    assert xm.xla_device_hw(torch_xla.device()) == 'TPU'
 
   def test_single_process(self):
     with concurrent.futures.ProcessPoolExecutor(max_workers=1) as executor:

test/pjrt/test_runtime_tpu.py

@@ -129,7 +129,7 @@ def test_xla_devices_single_process_one_chip_one_device_spawn(self):
 
   def test_default_xla_devices(self):
     with concurrent.futures.ProcessPoolExecutor(max_workers=1) as e:
-      f = e.submit(xm.get_xla_supported_devices, 'TPU')
+      f = e.submit(xm.get_xla_supported_devices)
       devices = [torch.device(d) for d in f.result()]
 
     self.assertListEqual(

test/spmd/test_xla_spmd_python_api_interaction.py

@@ -22,8 +22,7 @@ def setUpClass(cls):
     super().setUpClass()
 
   def test_get_xla_supported_devices(self):
-    device_type = os.environ['PJRT_DEVICE']
-    devices = xm.get_xla_supported_devices(device_type)
+    devices = xm.get_xla_supported_devices()
     self.assertEqual(len(devices), 1)
 
   def test_world_size(self):

test/test_autocast.py

@@ -348,7 +348,8 @@ def compare(first, second):
     self.assertFalse(self.is_autocast_enabled())
 
 
-@unittest.skipIf(not xm.get_xla_supported_devices("TPU"), f"TPU autocast test.")
+@unittest.skipIf(
+    xm.xla_device_hw(torch_xla.device()) != 'TPU', f"TPU autocast test.")
 class TestAutocastTPU(TestAutocastBase):
 
   @classmethod
@@ -404,7 +405,7 @@ class TestOtherOps(unittest.TestCase):
 
   # On TPU, the input of batch norm is casted into fp32, see torch_xla/csrc/autocast_mode.cpp
   @unittest.skipIf(
-      not xm.get_xla_supported_devices("TPU"),
+      xm.xla_device_hw(torch_xla.device()) != 'TPU',
       "the behavior of batch_norm autocast on TPU is different from others")
   def test_batch_norm_tpu(self):
     device = torch_xla.device()

torch_xla/core/xla_model.py

@@ -66,38 +66,21 @@ def is_xla_tensor(tensor: torch.Tensor) -> bool:
   return tensor.device.type == 'xla'
 
 
-def get_xla_supported_devices(devkind: Optional[str] = None,
-                              max_devices: Optional[int] = None) -> List[str]:
+def get_xla_supported_devices(max_devices: Optional[int] = None) -> List[str]:
   """Returns a list of supported devices of a given kind.
 
   Args:
-    devkind (string..., optional): If specified, a device type such as `TPU`,
-      `CUDA`, `CPU`, or name of custom PJRT device.
     max_devices (int, optional): The maximum number of devices to be returned of
       that kind.
 
   Returns:
     The list of device strings such as ['xla:0', 'xla:1', ...]
   """
-  # TODO(wcromar): Remove `devkind` after 2.3 release cut. We no longer support
-  # multiple device types.
-  if not devkind:
-    devices = torch_xla._XLAC._xla_get_devices()
-    return [
-        f'xla:{i}'
-        for i, _ in enumerate(devices[:max_devices] if max_devices else devices)
-    ]
-  else:
-    warnings.warn("`devkind` argument is deprecated and will be removed in a "
-                  "future release.")
-
-  xla_devices = _DEVICES.value
-  kind_devices = []
-  for i, device in enumerate(xla_devices):
-    if re.match(devkind + r':\d+$', device):
-      kind_devices.append('xla:{}'.format(i))
-  if kind_devices:
-    return kind_devices[:max_devices] if max_devices else kind_devices
+  devices = torch_xla._XLAC._xla_get_devices()
+  return [
+      f'xla:{i}'
+      for i, _ in enumerate(devices[:max_devices] if max_devices else devices)
+  ]
 
 
 def get_local_ordinal() -> int: