Revert "Fix psum & psum_scatter sharding logic"

This reverts commit 0c848de.

Author: Utkarsh Sharma

src/benchmark_collectives.py

@@ -78,7 +78,7 @@ def psum_benchmark(
     # DCN benchmark
     if dcn_size > 1:
 
-        @partial(shard_map, mesh=mesh, in_specs=P("dcn", None), out_specs=P(None, None))
+        @partial(shard_map, mesh=mesh, in_specs=P("dcn", None), out_specs=P(None))
         def f(x):
             return jax.lax.psum(x, "dcn")
 
@@ -99,12 +99,12 @@ def f(x):
     # ICI benchmark
     if ici_size > 1:
 
-        @partial(shard_map, mesh=mesh, in_specs=P(None, "ici"), out_specs=P(None, None))
+        @partial(shard_map, mesh=mesh, in_specs=P(None, None), out_specs=P(None, None))
         def f(x):
             return jax.lax.psum(x, "ici")
 
         sharded_matrix = jax.device_put(
-            matrix, jax.sharding.NamedSharding(mesh, P(None, "ici"))
+            matrix, jax.sharding.NamedSharding(mesh, P(None, None))
         )
         jitted_op = jax.jit(f)
         ici_average_time_ms_list = simple_timeit(
@@ -235,12 +235,12 @@ def f(x):
     # ICI benchmark
     if ici_size > 1:
 
-        @partial(shard_map, mesh=mesh, in_specs=P(None, "ici"), out_specs=P(None, "ici"))
+        @partial(shard_map, mesh=mesh, in_specs=P(None, None), out_specs=P(None, "ici"))
         def f(x):
             return jax.lax.psum_scatter(x, "ici", tiled=True)
 
         sharded_matrix = jax.device_put(
-            matrix, jax.sharding.NamedSharding(mesh, P(None, "ici"))
+            matrix, jax.sharding.NamedSharding(mesh, P(None, None))
         )
         jitted_op = jax.jit(f)
         ici_average_time_ms_list = simple_timeit(
@@ -443,9 +443,9 @@ def all_gather_benchmark_calculate_metrics(
         # each sharded matrix size is matrix_size_gbyte / ici_size and then it needs
         # to use (ici_size - 1) steps in a ring algorithm
         ici_bandwidth_gbyte_s_list = [
-                matrix_size_gbyte
-                * (ici_size - 1)
-                / ici_size
+                matrix_size_gbyte 
+                * (ici_size - 1) 
+                / ici_size 
                 / (ici_average_time_ms / 1e3)
                 for ici_average_time_ms in ici_average_time_ms_list
         ]