Merge pull request #635 from google:mattdavidow-numpy-prod

PiperOrigin-RevId: 631470826

MaxText/max_utils.py

@@ -289,7 +289,7 @@ def fill_unspecified_mesh_axes(parallelism_vals, target_product, parallelism_typ
     ), f"Found unspecified values (-1) for more than one {parallelism_type}\
       parallelism axis. At most one axis can be unspecified."
 
-    determined_val = target_product / np.product(parallelism_vals) * -1
+    determined_val = target_product / np.prod(parallelism_vals) * -1
 
     assert (
         determined_val >= 1 and determined_val.is_integer
@@ -301,9 +301,9 @@ def fill_unspecified_mesh_axes(parallelism_vals, target_product, parallelism_typ
   target_type = "slices" if parallelism_type == "DCN" else "devices per slice"
 
   assert (
-      np.product(parallelism_vals) == target_product
+      np.prod(parallelism_vals) == target_product
   ), f"Number of {target_type} {target_product} does not match\
-    the product of the {parallelism_type} parallelism {np.product(parallelism_vals)}"
+    the product of the {parallelism_type} parallelism {np.prod(parallelism_vals)}"
 
   return parallelism_vals
 

pedagogical_examples/shardings.py

@@ -127,9 +127,9 @@ def simple_timeit(f, tries=5, verbose=True):
 
   # Assert that we have correct inputs of sharding that fit the number of chips
   assert (
-      np.product(dcn_parallelism) * np.product(ici_parallelism) == num_devices
+      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)}"
 
   multi_slice_env = hasattr(jax.devices()[0], "slice_index")
   # Create device mesh