[backport] Use np.prod instead of np.product (#7301) (#7310)

Co-authored-by: JackCaoG <59073027+JackCaoG@users.noreply.github.com>
pytorch · Jun 18, 2024 · a901eb8 · a901eb8
1 parent 7296133
commit a901eb8
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/test/spmd/test_sharding_strategies.py b/test/spmd/test_sharding_strategies.py
@@ -67,9 +67,9 @@
 
 num_devices = xr.global_runtime_device_count()
 
-assert np.product(dcn_parallelism) * np.product(
+assert np.prod(dcn_parallelism) * np.prod(
     ici_parallelism) == num_devices, f"Number of devices {num_devices} \
-    does not match the product of the parallelism {np.product(dcn_parallelism) * np.product(ici_parallelism)}"
+    does not match the product of the parallelism {np.prod(dcn_parallelism) * np.prod(ici_parallelism)}"
 
 # Use HybridMesh to optimize multislice topology
 mesh = xs.HybridMesh(

diff --git a/torch_xla/distributed/spmd/xla_sharding.py b/torch_xla/distributed/spmd/xla_sharding.py
@@ -268,7 +268,7 @@ def _create_device_mesh_for_nd_torus(
         indices = itertools.combinations(
             range(len(assignable_physical_mesh)), num_axes)
         for c_axes, c_indices in zip(axes, indices):
-          if np.product(c_axes) == logical_axis_size:
+          if np.prod(c_axes) == logical_axis_size:
             assignment[logical_axis_index] = c_indices
             # Zero the assigned physical axes.
             assignable_physical_mesh = [