Migrate runtime.xla_device in favor of core.xla_model.xla_device

Author: ghpvnist

test/spmd/test_xla_sharding.py

@@ -1484,10 +1484,10 @@ def test_xla_patched_linear(self):
     """
 
     from torch_xla.distributed.spmd.xla_sharding import XLAPatchedLinear
-    import torch_xla.runtime
+    import torch_xla.core.xla_model as xm
     import torch.nn.functional as F
 
-    with torch_xla.runtime.xla_device():
+    with xm.xla_device():
       torch_xla.manual_seed(42)
       x0 = torch.randn(2, 3, requires_grad=True)
       w0 = torch.randn(4, 3, requires_grad=True)

torch_xla/core/xla_model.py

@@ -142,12 +142,12 @@ def xla_device(n: Optional[int] = None,
   Args:
     n (int, optional): The specific instance (ordinal) to be returned. If
       specified, the specific XLA device instance will be returned. Otherwise
-      the first device of `devkind` will be returned.
+      the first device (default 0) will be returned.
     devkind (string..., optional): If specified, device type such as `TPU`,
       `CUDA`, `CPU`, or custom PJRT device. Deprecated.
 
   Returns:
-    A `torch.device` with the requested instance.
+    A `torch.device` with the requested instance of an XLA device.
   """
   # When SPMD is enabled, we always return `xla:0` to the user, and
   # under the hood we use virtual device logic for every xla tensor
@@ -156,7 +156,16 @@ def xla_device(n: Optional[int] = None,
     torch_xla._XLAC._xla_set_default_device(device)
     return torch.device(device)
 
-  return runtime.xla_device(n, devkind)
+  if n is None:
+    return torch.device(torch_xla._XLAC._xla_get_default_device())
+
+  devices = xm.get_xla_supported_devices(devkind=devkind)
+  if n > len(devices):
+    raise IndexError('Device index {} out of range in {}'.format(n, devices))
+
+  device = devices[n]
+  torch_xla._XLAC._xla_set_default_device(device)
+  return torch.device(device)
 
 
 def _xla_real_device(device: torch.device) -> Any:

torch_xla/experimental/scan.py

@@ -120,7 +120,7 @@ def scan(fn, init, xs):
   Example:
 
     >>> # Example of using `scan` to implement `torch.cumsum`.
-    >>> import torch_xla.runtime
+    >>> import torch_xla.core.xla_model as xm
     >>> import torch
     >>> from torch_xla.experimental.scan import scan
     >>>
@@ -129,7 +129,7 @@ def scan(fn, init, xs):
     >>>   y = new_carry
     >>>   return new_carry, y
     >>>
-    >>> with torch_xla.runtime.xla_device():
+    >>> with xm.xla_device():
     >>>   init = torch.tensor([0.0, 0.0], requires_grad=True)
     >>>   xs = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]],
     >>>                     requires_grad=True)

torch_xla/experimental/scan_layers.py

@@ -47,11 +47,11 @@ def scan_layers(layers: Iterable[torch.nn.Module],
 
   Example:
 
-    >>> import torch_xla.runtime
+    >>> import torch_xla.core.xla_model as xm
     >>> import torch
     >>> import torch.nn as nn
     >>> from torch_xla.experimental.scan_layers import scan_layers
-    >>> with torch_xla.runtime.xla_device():
+    >>> with xm.xla_device():
     >>>   layers = [nn.Linear(16, 16) for i in range(10)]
     >>>   input = torch.randn(16)
     >>> output = scan_layers(layers, input)

torch_xla/runtime.py

@@ -103,30 +103,6 @@ def is_bf16_supported():
     return False
 
 
-def xla_device(n: Optional[int] = None,
-               devkind: Optional[str] = None) -> torch.device:
-  """Returns an XLA device.
-
-  Args:
-    n: Index of XLA device within visibible devices. If not set, use local
-      ordinal (default 0) to select an addressable device.
-    devkind: Type of device to return. Should match `device_type()`.
-
-  Returns:
-    A `torch.device` representing an XLA device.
-  """
-  if n is None:
-    return torch.device(torch_xla._XLAC._xla_get_default_device())
-
-  devices = xm.get_xla_supported_devices(devkind=devkind)
-  if n > len(devices):
-    raise IndexError('Device index {} out of range in {}'.format(n, devices))
-
-  device = devices[n]
-  torch_xla._XLAC._xla_set_default_device(device)
-  return torch.device(device)
-
-
 def local_process_count() -> int:
   """Returns the number of processes running on this host."""
   return xu.getenv_as(xenv.PJRT_LOCAL_PROCESS_COUNT, int, defval=1)
@@ -180,7 +156,7 @@ def local_ordinal() -> int:
   Local ordinal is in range [0, local_device_count)."""
   local_rank = xu.getenv_as(xenv.PJRT_LOCAL_PROCESS_RANK, int, 0)
   devices_per_process = addressable_device_count()
-  return local_rank * devices_per_process + xla_device().index
+  return local_rank * devices_per_process + xm.xla_device().index
 
 
 def process_index() -> int: