Improve PP->SS cost (#214)

Before we were only considering the PP->S(0)S(0) case, but the same logic applies to all shardings. We should generalize this to take the smallest cost into account, irrespective of mesh dims, but we should do this is a way that is efficient

Author: fmassa

autoparallel/collective_runtime_estimation.py

@@ -138,9 +138,8 @@ def redistribute_cost(
 
 def estimate_strategy_comms_cost(src_spec, tgt_spec):
     order = list(range(src_spec.mesh.ndim))
-    if src_spec.placements == (Partial(), Partial()) and tgt_spec.placements == (
-        Shard(0),
-        Shard(0),
+    if src_spec.placements == (Partial(), Partial()) and all(
+        p.is_shard() for p in tgt_spec.placements
     ):
         order = [1, 0]
     comms_cost = redistribute_cost(src_spec, tgt_spec, order)