Move spmd out of experimental

pytorch · Nov 10, 2023 · 558a894 · 558a894
1 parent 2ac9dc0
commit 558a894
Show file tree

Hide file tree

Showing 9 changed files with 11 additions and 11 deletions.
diff --git a/docs/spmd.md b/docs/spmd.md
@@ -46,8 +46,8 @@ import numpy as np
 import torch
 import torch_xla.core.xla_model as xm
 import torch_xla.runtime as xr
-import torch_xla.experimental.spmd.xla_sharding as xs
-from torch_xla.experimental.spmd import Mesh
+import torch_xla.distributed.spmd.xla_sharding as xs
+from torch_xla.distributed.spmd import Mesh
 
 # Enable XLA SPMD execution mode.
 xr.use_spmd()
@@ -104,7 +104,7 @@ We abstract logical mesh with [Mesh API](https://github.com/pytorch/xla/blob/028
 
 ```python
 import torch_xla.runtime as xr
-from torch_xla.experimental.spmd import Mesh
+from torch_xla.distributed.spmd import Mesh
 
 # Assuming you are running on a TPU host that has 8 devices attached
 num_devices = xr.global_runtime_device_count()

diff --git a/test/test_operations.py b/test/test_operations.py
@@ -36,7 +36,7 @@
 import torch_xla.debug.metrics as met
 import torch_xla.debug.model_comparator as mc
 import torch_xla.distributed.parallel_loader as pl
-import torch_xla.experimental.spmd.xla_sharding as xs
+import torch_xla.distributed.spmd.xla_sharding as xs
 from torch_xla import runtime as xr
 import torch_xla.test.test_utils as xtu
 import torch_xla.utils.utils as xu

diff --git a/torch_xla/experimental/spmd/__init__.py → torch_xla/distributed/spmd/__init__.py b/torch_xla/experimental/spmd/__init__.py → torch_xla/distributed/spmd/__init__.py
diff --git a/torch_xla/experimental/spmd/api.py → torch_xla/distributed/spmd/api.py b/torch_xla/experimental/spmd/api.py → torch_xla/distributed/spmd/api.py
diff --git a/...a/experimental/spmd/xla_sharded_tensor.py → ...la/distributed/spmd/xla_sharded_tensor.py b/...a/experimental/spmd/xla_sharded_tensor.py → ...la/distributed/spmd/xla_sharded_tensor.py
@@ -137,7 +137,7 @@ def sharding_spec(self):
 
   @property
   def sharding_type(self) -> 'ShardingType':
-    from torch_xla.experimental.spmd import ShardingType
+    from torch_xla.distributed.spmd import ShardingType
     sharding_type = torch_xla._XLAC._get_xla_sharding_type(self.global_tensor)
     return ShardingType(sharding_type)
 

diff --git a/torch_xla/experimental/spmd/xla_sharding.py → torch_xla/distributed/spmd/xla_sharding.py b/torch_xla/experimental/spmd/xla_sharding.py → torch_xla/distributed/spmd/xla_sharding.py
diff --git a/torch_xla/experimental/distributed_checkpoint/_helpers.py b/torch_xla/experimental/distributed_checkpoint/_helpers.py
@@ -5,7 +5,7 @@
 import dataclasses
 
 import torch
-import torch_xla.experimental.spmd.xla_sharding as xs
+import torch_xla.distributed.spmd.xla_sharding as xs
 
 from torch.distributed.checkpoint.planner import SavePlan
 from typing import (
@@ -23,7 +23,7 @@
 )
 from torch.distributed.checkpoint.metadata import (MetadataIndex,
                                                    STATE_DICT_TYPE)
-from torch_xla.experimental.spmd import XLAShardedTensor, ShardingType
+from torch_xla.distributed.spmd import XLAShardedTensor, ShardingType
 from torch.utils._pytree import tree_map
 
 PATH_ITEM = Union[str, int]

diff --git a/torch_xla/experimental/xla_sharded_tensor.py b/torch_xla/experimental/xla_sharded_tensor.py
@@ -4,7 +4,7 @@
 
 warnings.warn(
     "Importing from `torch_xla.experimental.xla_sharded_tensor` will be deprecated "
-    "after 2.2 release. Please use `torch_xla.experimental.spmd` "
+    "after 2.2 release. Please use `torch_xla.distributed.spmd` "
     "instead.", DeprecationWarning, 2)
 
-from .spmd.xla_sharded_tensor import *
+from torch_xla.distributed.spmd.xla_sharded_tensor import *
diff --git a/torch_xla/experimental/xla_sharding.py b/torch_xla/experimental/xla_sharding.py
@@ -4,7 +4,7 @@
 
 warnings.warn(
     "Importing from `torch_xla.experimental.xla_sharding` will be deprecated "
-    "after 2.2 release. Please use `torch_xla.experimental.spmd` instead.",
+    "after 2.2 release. Please use `torch_xla.distributed.spmd` instead.",
     DeprecationWarning, 2)
 
-from .spmd.xla_sharding import *
+from torch_xla.distributed.spmd.xla_sharding import *