fix circular import

torch/distributed/_tensor/_utils.py

@@ -123,13 +123,3 @@ def compute_global_tensor_info(
         elif not isinstance(placement, (Replicate, _Partial)):
             raise RuntimeError(f"placement type {type(placement)} not supported!")
     return tensor_shape, tensor_stride
-
-
-def _get_device_handle(device_type: str = "cuda"):
-    """
-    Get the module corresponding to the device_type which is cuda or cuda-like device.
-    For example, when the device_type is cuda, the module `torch.cuda` is returned.
-    Return None when device_type is cpu or there is no corresponding module,
-    otherwise return the corresponding module.
-    """
-    return getattr(torch, device_type, None) if device_type != "cpu" else None

torch/distributed/_tensor/device_mesh.py

@@ -5,7 +5,6 @@
 import torch
 import torch.distributed._functional_collectives as funcol
 
-from torch.distributed._tensor._utils import _get_device_handle
 from torch.distributed.distributed_c10d import (
     _get_default_group,
     all_gather,
@@ -50,6 +49,16 @@ def get_current_mesh(self) -> "DeviceMesh":
 mesh_resources: _MeshEnv = _MeshEnv()
 
 
+def _get_device_handle(device_type: str = "cuda"):
+    """
+    Get the module corresponding to the device_type which is cuda or cuda-like device.
+    For example, when the device_type is cuda, the module `torch.cuda` is returned.
+    Return None when device_type is cpu or there is no corresponding module,
+    otherwise return the corresponding module.
+    """
+    return getattr(torch, device_type, None) if device_type != "cpu" else None
+
+
 class DeviceMesh(object):
     """
     DeviceMesh represents a mesh of devices, where layout of devices could be

torch/distributed/_tensor/random.py

@@ -7,8 +7,7 @@
 import torch.distributed as dist
 
 from torch import Tensor
-from torch.distributed._tensor._utils import _get_device_handle
-from torch.distributed._tensor.device_mesh import DeviceMesh
+from torch.distributed._tensor.device_mesh import DeviceMesh, _get_device_handle
 from torch.distributed._tensor.placement_types import DTensorSpec, Shard