Modify sbp.split()'s karg: axis to dim

oneflow/api/common/sbp.h

@@ -33,7 +33,7 @@ inline Maybe<std::string> SbpToString(Symbol<SbpParallel> sbp_sym) {
   } else if (sbp_sym->has_partial_sum_parallel()) {
     sbp_str += "partial_sum";
   } else if (sbp_sym->has_split_parallel()) {
-    sbp_str += "split(axis=" + std::to_string(sbp_sym->split_parallel().axis()) + ")";
+    sbp_str += "split(dim=" + std::to_string(sbp_sym->split_parallel().axis()) + ")";
   } else {
     UNIMPLEMENTED_THEN_RETURN();
   }

python/oneflow/framework/distribute.py

@@ -14,16 +14,17 @@
 limitations under the License.
 """
 import traceback
+import warnings
 from contextlib import contextmanager
 
 import oneflow._oneflow_internal
 
 
-def split_sbp(axis: int) -> oneflow._oneflow_internal.sbp.sbp:
-    """Generate a split scheme in which op will be splitted at `axis`.
+def split_sbp(*dim, **kwarg) -> oneflow._oneflow_internal.sbp.sbp:
+    """Generate a split scheme in which op will be splitted at `dim`.
 
     Args:
-        axis (int): At `axis` the op will be splitted.
+        dim (int): At `dim` the op will be splitted.
 
     Returns:
         SbpParallel: Split scheme object, often required by `to_global` method of `Tensor`
@@ -34,5 +35,20 @@ def split_sbp(axis: int) -> oneflow._oneflow_internal.sbp.sbp:
         ct2 = t1.to_global(sbp=flow.sbp.split(0), placement=("cuda", ranks=[0, 1, 2, 3]))
 
     """
-    assert type(axis) is int
-    return oneflow._oneflow_internal.sbp.split(axis)
+    if len(kwarg) != 0:
+        if "dim" in kwarg:
+            dim = kwarg["dim"]
+        elif "axis" in kwarg:
+            dim = kwarg["axis"]
+            warnings.warn(
+                "This 'axis' parameter of oneflow.sbp.split() has been updated to 'dim' since OneFlow version 0.8."
+            )
+        else:
+            raise ValueError("Invalid keyword arguments, expected to be 'dim'.")
+
+    else:
+        assert len(dim) == 1
+        dim = dim[0]
+
+    assert type(dim) is int
+    return oneflow._oneflow_internal.sbp.split(dim)

python/oneflow/framework/docstr/tensor.py

@@ -319,13 +319,13 @@
 
         >>> # results on rank 0
         oneflow.Size([4])
-        tensor([0., 1., 0., 1.], placement=oneflow.placement(type="cpu", ranks=[0, 1]), sbp=(oneflow.sbp.split(axis=0),), dtype=oneflow.float32) 
+        tensor([0., 1., 0., 1.], placement=oneflow.placement(type="cpu", ranks=[0, 1]), sbp=(oneflow.sbp.split(dim=0),), dtype=oneflow.float32) 
 
     .. code-block:: python
 
         >>> # results on rank 1
         oneflow.Size([4])
-        tensor([0., 1., 0., 1.], placement=oneflow.placement(type="cpu", ranks=[0, 1]), sbp=(oneflow.sbp.split(axis=0),), dtype=oneflow.float32)
+        tensor([0., 1., 0., 1.], placement=oneflow.placement(type="cpu", ranks=[0, 1]), sbp=(oneflow.sbp.split(dim=0),), dtype=oneflow.float32)
     """,
 )
 
@@ -365,13 +365,13 @@
 
         >>> # results on rank 0
         oneflow.Size([2])
-        tensor([0., 1.], placement=oneflow.placement(type="cpu", ranks=[0, 1]), sbp=(oneflow.sbp.split(axis=0),), dtype=oneflow.float32)
+        tensor([0., 1.], placement=oneflow.placement(type="cpu", ranks=[0, 1]), sbp=(oneflow.sbp.split(dim=0),), dtype=oneflow.float32)
 
     .. code-block:: python
 
         >>> # results on rank 1
         oneflow.Size([2])
-        tensor([0., 1.], placement=oneflow.placement(type="cpu", ranks=[0, 1]), sbp=(oneflow.sbp.split(axis=0),), dtype=oneflow.float32)
+        tensor([0., 1.], placement=oneflow.placement(type="cpu", ranks=[0, 1]), sbp=(oneflow.sbp.split(dim=0),), dtype=oneflow.float32)
     """,
 )
 
@@ -424,13 +424,13 @@
 
         >>> # results on rank 0
         oneflow.Size([4])
-        tensor([0., 1., 0., 1.], placement=oneflow.placement(type="cpu", ranks=[0, 1]), sbp=(oneflow.sbp.split(axis=0),), dtype=oneflow.float32) 
+        tensor([0., 1., 0., 1.], placement=oneflow.placement(type="cpu", ranks=[0, 1]), sbp=(oneflow.sbp.split(dim=0),), dtype=oneflow.float32) 
 
     .. code-block:: python
 
         >>> # results on rank 1
         oneflow.Size([4])
-        tensor([0., 1., 0., 1.], placement=oneflow.placement(type="cpu", ranks=[0, 1]), sbp=(oneflow.sbp.split(axis=0),), dtype=oneflow.float32)
+        tensor([0., 1., 0., 1.], placement=oneflow.placement(type="cpu", ranks=[0, 1]), sbp=(oneflow.sbp.split(dim=0),), dtype=oneflow.float32)
 
     For global tensor:
 
@@ -447,13 +447,13 @@
 
         >>> # results on rank 0
         oneflow.Size([2])
-        tensor([0., 1.], placement=oneflow.placement(type="cpu", ranks=[0, 1]), sbp=(oneflow.sbp.split(axis=0),), dtype=oneflow.float32)
+        tensor([0., 1.], placement=oneflow.placement(type="cpu", ranks=[0, 1]), sbp=(oneflow.sbp.split(dim=0),), dtype=oneflow.float32)
 
     .. code-block:: python
 
         >>> # results on rank 1
         oneflow.Size([2])
-        tensor([0., 1.], placement=oneflow.placement(type="cpu", ranks=[0, 1]), sbp=(oneflow.sbp.split(axis=0),), dtype=oneflow.float32)
+        tensor([0., 1.], placement=oneflow.placement(type="cpu", ranks=[0, 1]), sbp=(oneflow.sbp.split(dim=0),), dtype=oneflow.float32)
     """,
 )
 

python/oneflow/framework/docstr/tensor_attributes.py

@@ -97,9 +97,9 @@
 
     ``oneflow.sbp`` includes three types:
 
-        - oneflow.sbp.split(axis)
+        - oneflow.sbp.split(dim)
 
-          Indicates that the global tensor is evenly divided according to the dimension `axis` and distributed on each rank.
+          Indicates that the global tensor is evenly divided according to the dimension `dim` and distributed on each rank.
 
         - oneflow.sbp.broadcast()
 
@@ -120,7 +120,7 @@
 
         >>> s = flow.sbp.split(0)
         >>> s
-        oneflow.sbp.split(axis=0)
+        oneflow.sbp.split(dim=0)
         >>> b = flow.sbp.broadcast()
         >>> b
         oneflow.sbp.broadcast

python/oneflow/test/graph/test_graph_zero.py

@@ -206,7 +206,7 @@ def one_eval_iter():
 
         for state in linear_t_g._state():
             test_case.assertEqual(
-                state.origin.sbp, (oneflow.sbp.split(axis=0), oneflow.sbp.split(axis=0))
+                state.origin.sbp, (oneflow.sbp.split(dim=0), oneflow.sbp.split(dim=0))
             )
 
         # In evaluation graph, paramters's sbp are flow.sbp.split(0).