Address comments -- fix typos and variable names

pytorch · Nov 7, 2023 · 1808663 · 1808663
1 parent a98bfb2
commit 1808663
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 14 deletions.
diff --git a/torch_xla/csrc/init_python_bindings.cpp b/torch_xla/csrc/init_python_bindings.cpp
@@ -1698,10 +1698,10 @@ void InitXlaModuleBindings(py::module m) {
         [](const at::Tensor& input, const py::list& tile_assignment,
            const py::list& group_assignment, const py::list& replication_groups,
            int sharding_type) {
-          c10::List<at::IntArrayRef> time_assignment_list =
+          c10::List<at::IntArrayRef> tile_assignment_list =
               c10::List<at::IntArrayRef>();
           for (auto t : tile_assignment) {
-            time_assignment_list.push_back(
+            tile_assignment_list.push_back(
                 at::IntArrayRef(t.cast<std::vector<int64_t>>()));
           }
 
@@ -1720,7 +1720,7 @@ void InitXlaModuleBindings(py::module m) {
           }
 
           xla_mark_sharding_dynamo_custom_op(
-              input, time_assignment_list, group_assignment_list,
+              input, tile_assignment_list, group_assignment_list,
               replication_groups_list, sharding_type);
         });
   m.def("_xla_clear_sharding", [](const at::Tensor& input) {

diff --git a/torch_xla/experimental/xla_sharding.py b/torch_xla/experimental/xla_sharding.py
@@ -83,7 +83,7 @@ def get_axis_name_idx(self, name: str) -> int:
   @functools.lru_cache(maxsize=None)
   def get_op_sharding(self,
                       partition_spec: Tuple,
-                      flatten=False) -> torch_xla._XLAC.OpSharding:
+                      flatten_opsharding = False) -> torch_xla._XLAC.OpSharding:
     """
     Return the OpSharding for the given partition spec. This is an expensive
     operation as the mesh grows, so the value is cached for reuse.
@@ -107,7 +107,7 @@ def get_op_sharding(self,
         sharding_type, tile_assignment, len(partition_spec), replicate_dims)
 
     # If flatten = True, return the flattened version of OpSharding
-    if flatten:
+    if flatten_opsharding:
       return (tile_assignment.tolist(), group_assignment, replication_groups,
               int(sharding_type))
     else:
@@ -459,7 +459,7 @@ def _translate_named_partition_spec(mesh: Mesh, partition_spec: Tuple):
 def mark_sharding(t: Union[torch.Tensor, XLAShardedTensor],
                   mesh: Mesh,
                   partition_spec: Tuple[Union[Tuple, int, str, None]],
-                  dynamo_custom_op: bool = False) -> XLAShardedTensor:
+                  use_dynamo_custom_op: bool = False) -> XLAShardedTensor:
   """
     Annotates the tensor provided with XLA partition spec. Internally,
     it annotates the corresponding XLATensor as sharded for the XLA SpmdPartitioner pass.
@@ -508,7 +508,7 @@ def mark_sharding(t: Union[torch.Tensor, XLAShardedTensor],
   assert len(t.shape) == len(partition_spec), \
     f"Partition spec length ({len(partition_spec)}) should be equal to the input rank ({len(t.shape)})."
 
-  if dynamo_custom_op:
+  if use_dynamo_custom_op:
     tile_assignment, group_assignment, replication_groups, sharding_type = mesh.get_op_sharding(
         partition_spec, flatten=True)
 
@@ -517,19 +517,21 @@ def mark_sharding(t: Union[torch.Tensor, XLAShardedTensor],
           t.global_tensor, tile_assignment, group_assignment,
           replication_groups, sharding_type)
       return t
-    torch_xla._XLAC._xla_mark_sharding_dynamo_custom_op(t, tile_assignment,
-                                                        group_assignment,
-                                                        replication_groups,
-                                                        sharding_type)
-    return XLAShardedTensor(t)
+    else:
+      torch_xla._XLAC._xla_mark_sharding_dynamo_custom_op(t, tile_assignment,
+                                                          group_assignment,
+                                                          replication_groups,
+                                                          sharding_type)
+      return XLAShardedTensor(t)
   else:
     op_sharding = mesh.get_op_sharding(partition_spec)
 
     if isinstance(t, XLAShardedTensor):
       torch_xla._XLAC._xla_mark_sharding(t.global_tensor, op_sharding)
       return t
-    torch_xla._XLAC._xla_mark_sharding(t, op_sharding)
-    return XLAShardedTensor(t)
+    else:
+      torch_xla._XLAC._xla_mark_sharding(t, op_sharding)
+      return XLAShardedTensor(t)
 
 
 def clear_sharding(t: Union[torch.Tensor, XLAShardedTensor]) -> torch.Tensor: