Fix matrix_utils_test to be compatible with NumPy 2. Update dtype to …

…jnp.int32 where applicable to avoid overflow errors.

PiperOrigin-RevId: 673207362

jaxite/jaxite_lib/matrix_utils.py

@@ -58,7 +58,7 @@ def x_power_n_minus_1(n: jnp.uint32, poly_mod_deg: jnp.uint32) -> jnp.ndarray:
   degree = n % (2 * poly_mod_deg)
   flip = degree // poly_mod_deg
   reduced_degree = degree % poly_mod_deg
-  zeros = jnp.zeros(poly_mod_deg, dtype=jnp.uint32)
+  zeros = jnp.zeros(poly_mod_deg, dtype=jnp.int32)
   return zeros.at[reduced_degree].set((-1) ** flip) - zeros.at[0].set(1)
 
 

jaxite/jaxite_lib/matrix_utils_test.py

@@ -207,22 +207,22 @@ def test_monomial_mul_32_bit_modulus(self, degree, expected):
     np.testing.assert_array_equal(expected, mono_mul_output)
 
   def test_x_power_n_minus_1(self):
-    expected = jnp.array([-1, 0, 1, 0], dtype=jnp.uint32)
+    expected = jnp.array([-1, 0, 1, 0], dtype=jnp.int32)
     actual = matrix_utils.x_power_n_minus_1(n=2, poly_mod_deg=4)
     np.testing.assert_array_equal(expected, actual)
 
   def test_x_power_n_minus_1_zero(self):
-    expected = jnp.array([0, 0, 0, 0], dtype=jnp.uint32)
+    expected = jnp.array([0, 0, 0, 0], dtype=jnp.int32)
     actual = matrix_utils.x_power_n_minus_1(n=0, poly_mod_deg=4)
     np.testing.assert_array_equal(expected, actual)
 
   def test_x_power_n_minus_1_reduced_degree_with_sign_flip(self):
-    expected = jnp.array([-1, 0, -1, 0], dtype=jnp.uint32)
+    expected = jnp.array([-1, 0, -1, 0], dtype=jnp.int32)
     actual = matrix_utils.x_power_n_minus_1(n=6, poly_mod_deg=4)
     np.testing.assert_array_equal(expected, actual)
 
   def test_x_power_n_minus_1_reduced_degree_without_sign_flip(self):
-    expected = jnp.array([-1, 0, 1, 0], dtype=jnp.uint32)
+    expected = jnp.array([-1, 0, 1, 0], dtype=jnp.int32)
     actual = matrix_utils.x_power_n_minus_1(n=10, poly_mod_deg=4)
     np.testing.assert_array_equal(expected, actual)