Fix unspecified mesh-axis integer validation

Author: huytransformer

src/maxtext/utils/max_utils.py

@@ -455,7 +455,7 @@ def fill_unspecified_mesh_axes(parallelism_vals, target_product, parallelism_typ
     determined_val = target_product / np.prod(parallelism_vals) * -1
 
     assert (
-        determined_val >= 1 and determined_val.is_integer
+        determined_val >= 1 and determined_val.is_integer()
     ), f"Unspecified value unable to be determined with the given\
       {parallelism_type} parallelism values"
 

tests/unit/max_utils_test.py

@@ -137,6 +137,20 @@ def test_invalid_strategy(self):
       max_utils.is_valid_custom_mesh([1, 1, 1, 1, 1, 16, 16, 1], "invalid_strategy")
 
 
+class FillUnspecifiedMeshAxesTest(unittest.TestCase):
+  """Tests for fill_unspecified_mesh_axes."""
+
+  def test_rejects_non_integer_unspecified_value(self):
+    with self.assertRaises(AssertionError):
+      max_utils.fill_unspecified_mesh_axes([2, -1, 3], 10, "ICI")
+
+  def test_fills_integer_unspecified_value(self):
+    self.assertEqual(
+        max_utils.fill_unspecified_mesh_axes([2, -1, 5], 20, "ICI"),
+        [2, 2, 5],
+    )
+
+
 class UnscanTest(unittest.TestCase):
   """Test unscanning utility."""