Fix other callsites

Author: wconstab

autoparallel/compute_estimation.py

@@ -155,7 +155,7 @@ def _get_device_tflops(dtype):
     return device_limit.gemm_tflops[dtype]
 
 
-def _get_sharded_shape(spec):
+def _get_sharded_shape_stride(spec):
     mesh = spec.mesh
     tensor_shape = spec.tensor_meta.shape
     # TODO: take dtype into account as well
@@ -196,7 +196,7 @@ def estimate_strategy_runtime_cost(node, strategy):
         for k, v in kwargs.items():
             assert not isinstance(v, torch.Tensor), f"{node} {v}"
     args_sizes_strides = tuple(
-        _get_sharded_shape(spec) for spec in strategy.input_specs
+        _get_sharded_shape_stride(spec) for spec in strategy.input_specs
     )
 
     counter = 0

autoparallel/optimize_sharding.py

@@ -11,7 +11,10 @@
 from torch.distributed.tensor.placement_types import Replicate, Shard
 from torch.utils._pytree import tree_flatten, tree_map_only
 
-from .compute_estimation import _get_sharded_shape, estimate_strategy_runtime_cost
+from .compute_estimation import (
+    _get_sharded_shape_stride,
+    estimate_strategy_runtime_cost,
+)
 from .propagation_rules import _create_all_options
 from .utils import get_placement_options
 
@@ -488,7 +491,7 @@ def add_parameter_memory_constraint(self, memory_factor_low, memory_factor_high)
                 data = self.ds[(s_i, 0, ii, 0)]
                 spec = data["inp_strat"]
                 tensor_shape = spec.tensor_meta.shape
-                new_tensor_shape = _get_sharded_shape(spec)
+                new_tensor_shape, _ = _get_sharded_shape_stride(spec)
                 new_size = math.prod(new_tensor_shape)
                 old_size = math.prod(tensor_shape)
                 elms.append(data["va"] * new_size / old_size)