feat: NTT

algebra/ff/babybear.py

@@ -4,3 +4,4 @@
 class BabyBear(PrimeField):
   # BabyBear prime: 2^31 - 2^27 + 1
   P = 2013265921
+  w = 31

algebra/ff/m31.py

@@ -3,4 +3,5 @@
 
 class M31(PrimeField):
   # Mersenne31 prime: 2^31 - 1
-  P = 2147483647
+  P = 2**31 - 1
+  w = 7

algebra/ff/prime_field.py

@@ -4,6 +4,7 @@
 
 class PrimeField:
   P: int = None
+  w: int = None
 
   def __init__(self, x):
     if isinstance(x, (int, float, list, Tensor)):

algebra/poly/ntt.py

@@ -0,0 +1,122 @@
+from tinygrad.tensor import Tensor
+from tinygrad import dtypes
+import numpy as np
+
+
+def _next_valid_n(l, p):
+  """Find next n >= l that divides p-1 efficiently"""
+  p_minus_1 = p - 1
+  for n in range(l, min(p, 2 * l)):
+    if p_minus_1 % n == 0:
+      return n
+
+  bit_length = l.bit_length()
+  power_of_two = 1 << bit_length
+  if power_of_two < l:
+    power_of_two = power_of_two << 1
+
+  while power_of_two < p:
+    if p_minus_1 % power_of_two == 0:
+      return power_of_two
+    power_of_two = power_of_two << 1
+
+  for divisor in range(l, int(p_minus_1**0.5) + 1):
+    if p_minus_1 % divisor == 0:
+      quotient = p_minus_1 // divisor
+      if quotient >= l:
+        return quotient
+      return divisor
+
+  return p - 1
+
+
+def ntt(polynomial: Tensor, prime: int, primitive_root: int) -> Tensor:
+  """
+  Compute the Number Theoretic Transform of a polynomial.
+
+  Args:
+      polynomial: Coefficient vector of the polynomial (Tensor or list/array)
+      prime: The prime modulus
+      primitive_root: A primitive root modulo prime
+
+  Returns:
+      The NTT of the polynomial as a Tensor
+  """
+  return _transform(polynomial, prime, primitive_root, inverse=False)
+
+
+def intt(transformed: Tensor, prime: int, primitive_root: int) -> Tensor:
+  """
+  Compute the Inverse Number Theoretic Transform.
+
+  Args:
+      transformed: The transformed polynomial (Tensor or list/array)
+      prime: The prime modulus
+      primitive_root: A primitive root modulo prime
+
+  Returns:
+      The original polynomial coefficients as a Tensor
+  """
+  return _transform(transformed, prime, primitive_root, inverse=True)
+
+
+def _transform(x, prime: int, primitive_root: int, inverse: bool = False) -> Tensor:
+  """Internal function to perform NTT or INTT transformation"""
+  if not isinstance(x, Tensor):
+    x = Tensor(x, dtype=dtypes.uint64)
+
+  dtype = x.dtype
+  n = _next_valid_n(len(x), prime)
+  print(f"n: {n}, prime: {prime}, primitive_root: {primitive_root}")
+
+  if len(x) < n:
+    padded_np = Tensor.zeros(n, dtype=dtype).contiguous()
+    padded_np[: len(x)] = x
+    x = padded_np
+
+  # Create powers matrix (i*j) % n
+  i = Tensor.arange(n, dtype=dtypes.uint64).reshape(n, 1)
+  j = Tensor.arange(n, dtype=dtypes.uint64).reshape(1, n)
+  powers = (i * j) % n
+
+  # Get omega (or inverse omega for INTT)
+  omega = pow(primitive_root, (prime - 1) // n, prime)
+  if inverse:
+    omega = pow(omega, prime - 2, prime)  # Modular inverse
+
+  # Compute all powers of omega
+  omega_powers = np.ones(n, dtype=np.uint64)
+  current = 1
+  for k in range(n):
+    omega_powers[k] = current
+    current = (current * omega) % prime
+
+  # Create transformation matrix
+  matrix = Tensor(omega_powers, dtype=dtypes.uint64)[powers]
+  result = (matrix @ (x % prime)) % prime
+
+  # For inverse, apply scaling factor (1/n mod prime)
+  if inverse:
+    n_inv = pow(n, prime - 2, prime)
+    result = (result * n_inv) % prime
+
+  # Perform the transformation
+  return result
+
+
+if __name__ == "__main__":
+  from algebra.ff.m31 import M31
+  from random import randint
+
+  n, prime, primitive_root = 10, M31.P, M31.w
+  polynomial = Tensor([randint(0, prime - 1) for _ in range(n)], dtype=dtypes.uint64)
+  print(f"polynomial: {polynomial.numpy()}")
+
+  transformed = ntt(polynomial, prime, primitive_root)
+  print("NTT result:", transformed.numpy())
+
+  original = intt(transformed, prime, primitive_root)
+  print("INTT result:", original.numpy())
+
+  assert (polynomial.numpy() == original.numpy()[: len(polynomial)]).all(), "INTT(NTT(polynomial)) != polynomial"
+  print("Verification successful")

algebra/poly/univariate.py

@@ -1,6 +1,8 @@
 from tinygrad.tensor import Tensor
 from tinygrad import dtypes
 from algebra.ff.prime_field import PrimeField as PF
+from algebra.poly.ntt import ntt, intt
+import numpy as np
 
 
 class Polynomial:
@@ -14,17 +16,15 @@ class Polynomial:
     """
 
   def __init__(self, coeffs: list[int] | Tensor, prime_field: PF = None):
-    """
-    Initialize the polynomial.
-
-    coeffs: a list of field elements (instances of a PrimeField subclass)
-    prime_field: optional prime field class to use for modular arithmetic
-    """
+    """Initialize polynomial with coefficients and optional prime field."""
     self.PrimeField = prime_field
     if isinstance(coeffs, list):
+      coeffs = np.trim_zeros(coeffs, "b")
       self.coeffs = Tensor(coeffs, dtype=dtypes.int32)
-    elif isinstance(coeffs, Tensor):
-      self.coeffs = coeffs
+    else:
+      coeffs_np = coeffs.numpy()
+      coeffs_np = np.trim_zeros(coeffs_np, "b")
+      self.coeffs = Tensor(coeffs_np, dtype=coeffs.dtype)
 
   def degree(self) -> int:
     """
@@ -49,6 +49,20 @@ def evaluate(self, x: int | Tensor):
       result = result * x + coeff
     return result
 
+  def ntt(self):
+    """
+    Compute the Number Theoretic Transform of the polynomial.
+    """
+    p_ntt = ntt(self.coeffs.cast(dtypes.uint64), self.PrimeField.P, self.PrimeField.w).cast(self.coeffs.dtype)
+    return Polynomial(p_ntt, self.PrimeField)
+
+  def intt(self):
+    """
+    Compute the Inverse Number Theoretic Transform of the polynomial.
+    """
+    p_intt = intt(self.coeffs.cast(dtypes.uint64), self.PrimeField.P, self.PrimeField.w).cast(self.coeffs.dtype)
+    return Polynomial(p_intt, self.PrimeField)
+
   def __evaluate_all(self, xs: Tensor):
     """
     Evaluate the polynomial at all elements in xs using Horner's method.

pyproject.toml

@@ -8,6 +8,7 @@ authors = [{ name = "Pia", email = "gayeongparkk@gmail.com" }]
 requires-python = ">=3.10"
 dependencies = [
     "numpy>=2.2.3",
+    "sympy>=1.13.3",
     "tinygrad>=0.10.2",
 ]
 

tests/test_poly.py

@@ -2,6 +2,7 @@
 from algebra.ff.m31 import M31
 from algebra.ff.babybear import BabyBear
 from tinygrad import Tensor
+from random import randint
 
 
 def test_polynomial_operations_m31():
@@ -70,3 +71,8 @@ def test_polynomial_operations_babybear():
   # Test evaluate_all
   result = p1(Tensor([1, 2, 3])).numpy()
   assert (result == [6, 17, 34]).all()
+
+  p7 = Polynomial([randint(0, 100) for _ in range(8)], M31)
+  p7_ntt = p7.ntt()
+  p7_intt = p7_ntt.intt()
+  assert (p7_intt.coeffs.numpy() == p7.coeffs.numpy()).all()