internalize _parse_xla_device

torch_xla/_internal/utils.py

@@ -0,0 +1,8 @@
+import logging
+import re
+
+
+def parse_xla_device(device: str):
+  m = re.match(r'([A-Z]+):(\d+)$', device)
+  if m:
+    return (m.group(1), int(m.group(2)))

torch_xla/core/xla_model.py

@@ -19,6 +19,8 @@
 import torch_xla.utils.utils as xu
 import torch_xla.utils.closures as xc
 import os
+from torch_xla.experimental.deprecation import deprecated
+import torch_xla._internal.utils as _utils
 
 _DEVICES = xu.LazyProperty(lambda: torch_xla._XLAC._xla_get_devices())
 
@@ -40,6 +42,9 @@
 
 XLA_LIB = Library("xla", "DEF")
 
+from . import xla_model as this_module
+parse_xla_device = deprecated(this_module, _utils.parse_xla_device)
+
 
 def _init_world_size_ordinal():
   global _WORLD_SIZE, _ORDINAL
@@ -76,12 +81,6 @@ def is_xla_tensor(tensor):
   return tensor.device.type == 'xla'
 
 
-def parse_xla_device(device):
-  m = re.match(r'([A-Z]+):(\d+)$', device)
-  if m:
-    return (m.group(1), int(m.group(2)))
-
-
 def get_xla_supported_devices(devkind=None, max_devices=None):
   """Returns a list of supported devices of a given kind.
 

torch_xla/distributed/spmd/xla_sharding.py

@@ -4,6 +4,7 @@
 import torch
 import torch_xla
 import torch_xla.core.xla_model as xm
+import torch_xla._internal.utils as iutils
 from torch_xla.distributed.spmd import XLAShardedTensor, XLAShard
 import torch_xla.runtime as xr
 
@@ -219,7 +220,7 @@ def __init__(self,
     mesh_shape = tuple([x * y for x, y in zip(ici_mesh_shape, dcn_mesh_shape)])
     self.device_attributes = xr.global_runtime_device_attributes()
     self.device_attributes.sort(
-        key=lambda attr: xm.parse_xla_device(attr['name'])[1])
+        key=lambda attr: iutils.parse_xla_device(attr['name'])[1])
 
     if 'slice_index' in self.device_attributes[0] and np.prod(
         dcn_mesh_shape) == 1: