fix: rational dot product for 1-dim matrices (#27)

* fix: new way of doing matrix multiplication with rationals

* chore: reformated and added code comments

* fix: linting

Author: jcabrero

nada_algebra/array.py

@@ -19,6 +19,7 @@
     Integer,
     UnsignedInteger,
 )
+
 from nada_algebra.types import (
     Rational,
     SecretRational,
@@ -27,6 +28,9 @@
     secret_rational,
     get_log_scale,
 )
+
+from nada_algebra.context import UnsafeArithmeticSession
+
 from nada_algebra.utils import copy_metadata
 
 
@@ -270,59 +274,11 @@ def rational_matmul(self, other: "NadaArray") -> "NadaArray":
         Returns:
             NadaArray: A new NadaArray representing the result of matrix multiplication.
         """
-        return NadaArray(NadaArray.rational_matmul_recursive(self, other))
-
-    @staticmethod
-    def rational_matmul_recursive(A: "NadaArray", B: "NadaArray") -> "NadaArray":
-        """
-        Perform matrix multiplication with another NadaArray when both have Rational Numbers.
-        It improves the number of truncations to be needed to the resulting matrix dimensions mxp.
-
-        Args:
-            other (NadaArray): The NadaArray to perform matrix multiplication with.
-
-        Returns:
-            NadaArray: A new NadaArray representing the result of matrix multiplication.
-        """
-        # We check that both have same number of dimensions.
-        if A.ndim != B.ndim:
-            raise ValueError(
-                f"Matrices are not aligned for multiplication: {A.inner.shape} and {B.inner.shape}"
+        with UnsafeArithmeticSession():
+            return NadaArray(np.array(self.inner @ other.inner)).apply(
+                lambda x: x.rescale_down()
             )
 
-        # Since both have the same number of dimensions, we now check if they are 2D matrices.
-        # If they are not, they will pass this check and execute normally.
-        # Otherwise, we will do matrix contraction (i.e., compute matrix multiplication dimension by dimension).
-        if A.ndim > 2:
-            a = [
-                NadaArray.rational_matmul_recursive(A[i], B[i])
-                for i in range(A.shape[0])
-            ]  # We remove one dimension here.
-            return np.array(a)
-
-        # Get the dimensions of the matrices
-        (m, n) = A.shape
-        (n_, p) = B.shape
-
-        if n != n_:
-            raise ValueError(
-                f"Matrices are not aligned for multiplication: {A.inner.shape} and {B.inner.shape}"
-            )
-
-        # Initialize the result matrix C with zeros
-        C = np.zeros((m, p), dtype=object)
-
-        # Perform matrix multiplication
-        for i in range(m):
-            for j in range(p):
-                for k in range(n):
-                    if k == 0:
-                        C[i][j] = A[i][k].mul_no_rescale(B[k][j])
-                    else:
-                        C[i][j] += A[i][k].mul_no_rescale(B[k][j])
-                C[i][j] = C[i][j].rescale_down()
-        return C
-
     def __matmul__(self, other: Any) -> "NadaArray":
         """
         Perform matrix multiplication with another NadaArray.
@@ -357,6 +313,9 @@ def dot(self, other: "NadaArray") -> "NadaArray":
         Returns:
             NadaArray: A new NadaArray representing the dot product result.
         """
+        if self.is_rational or other.is_rational:
+            return self.rational_matmul(other)
+
         return NadaArray(self.inner.dot(other.inner))
 
     def hstack(self, other: "NadaArray") -> "NadaArray":

nada_algebra/context.py

@@ -0,0 +1,108 @@
+from nada_algebra.types import SecretRational, Rational, _NadaRational
+
+
+class UnsafeArithmeticSession:
+    """
+    A context manager that temporarily modifies the behavior of arithmetic operations
+    for Rational and SecretRational types, disabling rescaling for multiplication and division.
+
+    Attributes:
+        mul_rational (function): Original __mul__ method of Rational.
+        mul_secret_rational (function): Original __mul__ method of SecretRational.
+        truediv_rational (function): Original __truediv__ method of Rational.
+        truediv_secret_rational (function): Original __truediv__ method of SecretRational.
+    """
+
+    def __init__(self):
+        """
+        Initializes the UnsafeArithmeticSession by storing the original
+        multiplication and division methods of Rational and SecretRational.
+        """
+        self.mul_rational = Rational.__mul__
+        self.mul_secret_rational = SecretRational.__mul__
+
+        self.truediv_rational = Rational.__truediv__
+        self.truediv_secret_rational = SecretRational.__truediv__
+
+    def __enter__(self):
+        """
+        Enters the context, temporarily replacing the multiplication and division
+        methods of Rational and SecretRational to disable rescaling.
+        """
+
+        def mul_no_rescale_wrapper(self: Rational, other: _NadaRational):
+            """
+            Wrapper for Rational.__mul__ that disables rescaling.
+
+            Args:
+                self (Rational): The Rational instance.
+                other (_NadaRational): The other operand.
+
+            Returns:
+                Rational: Result of the multiplication without rescaling.
+            """
+            return Rational.mul_no_rescale(self, other, ignore_scale=True)
+
+        def secret_mul_no_rescale_wrapper(self: SecretRational, other: _NadaRational):
+            """
+            Wrapper for SecretRational.__mul__ that disables rescaling.
+
+            Args:
+                self (SecretRational): The SecretRational instance.
+                other (_NadaRational): The other operand.
+
+            Returns:
+                SecretRational: Result of the multiplication without rescaling.
+            """
+            return SecretRational.mul_no_rescale(self, other, ignore_scale=True)
+
+        def divide_no_rescale_wrapper(self: Rational, other: _NadaRational):
+            """
+            Wrapper for Rational.__truediv__ that disables rescaling.
+
+            Args:
+                self (Rational): The Rational instance.
+                other (_NadaRational): The other operand.
+
+            Returns:
+                Rational: Result of the division without rescaling.
+            """
+            return Rational.divide_no_rescale(self, other, ignore_scale=True)
+
+        def secret_divide_no_rescale_wrapper(
+            self: SecretRational, other: _NadaRational
+        ):
+            """
+            Wrapper for SecretRational.__truediv__ that disables rescaling.
+
+            Args:
+                self (SecretRational): The SecretRational instance.
+                other (_NadaRational): The other operand.
+
+            Returns:
+                SecretRational: Result of the division without rescaling.
+            """
+            return SecretRational.divide_no_rescale(self, other, ignore_scale=True)
+
+        Rational.__mul__ = mul_no_rescale_wrapper
+        SecretRational.__mul__ = secret_mul_no_rescale_wrapper
+        Rational.__truediv__ = divide_no_rescale_wrapper
+        SecretRational.__truediv__ = secret_divide_no_rescale_wrapper
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """
+        Exits the context, restoring the original multiplication and division methods
+        of Rational and SecretRational.
+
+        Args:
+            exc_type (type): Exception type if an exception occurred, else None.
+            exc_val (Exception): Exception instance if an exception occurred, else None.
+            exc_tb (traceback): Traceback object if an exception occurred, else None.
+        """
+        # Restore the original __mul__ method
+        Rational.__mul__ = self.mul_rational
+        SecretRational.__mul__ = self.mul_secret_rational
+
+        # Restore the original __truediv__ method
+        Rational.__truediv__ = self.truediv_rational
+        SecretRational.__truediv__ = self.truediv_secret_rational

nada_algebra/types.py

@@ -1,6 +1,7 @@
 """Additional special data types"""
 
 import warnings
+from functools import partial
 import numpy as np
 
 import nada_dsl as dsl

tests/nada-tests/nada-project.toml

@@ -152,4 +152,8 @@ prime_size = 128
 
 [[programs]]
 path = "src/matrix_multiplication_rational_multidim.py"
+prime_size = 128
+
+[[programs]]
+path = "src/dot_product_rational.py"
 prime_size = 128

tests/nada-tests/src/dot_product_rational.py

@@ -0,0 +1,25 @@
+from nada_dsl import *
+import nada_algebra as na
+
+
+def nada_main():
+    parties = na.parties(3)
+
+    a = na.array((3,), parties[0], "A", na.SecretRational)
+    b = na.array((3,), parties[1], "B", na.SecretRational)
+    c = na.ones((3,), na.Rational)
+
+    result = a.dot(b)
+
+    result_b = a @ b
+
+    result_c = a.dot(c)
+
+    result_d = a @ c
+
+    return (
+        result.output(parties[1], "my_output_a")
+        + result_b.output(parties[1], "my_output_b")
+        + result_c.output(parties[1], "my_output_c")
+        + result_d.output(parties[1], "my_output_d")
+    )

tests/nada-tests/src/matrix_multiplication_rational.py

@@ -1,21 +1,29 @@
 from nada_dsl import *
-
-# Step 0: Nada Algebra is imported with this line
 import nada_algebra as na
 
 
 def nada_main():
-    # Step 1: We use Nada Algebra wrapper to create "Party0", "Party1" and "Party2"
     parties = na.parties(3)
 
-    # Step 2: Party0 creates an array of dimension (3 x 3) with name "A"
     a = na.array([3, 3], parties[0], "A", na.SecretRational)
 
-    # Step 3: Party1 creates an array of dimension (3 x 3) with name "B"
     b = na.array([3, 3], parties[1], "B", na.SecretRational)
 
-    # Step 4: The result is of computing the dot product between the two which is another (3 x 3) matrix
-    result = a @ b
+    c = na.array((3,), parties[2], "C", na.SecretRational)
+
+    d = na.ones([3, 3], na.Rational)
+
+    result_a = a @ b
+
+    result_b = a @ c
+
+    result_c = a @ d
+
+    result_d = d @ a
 
-    # Step 5: We can use result.output() to produce the output for Party2 and variable name "my_output"
-    return result.output(parties[1], "my_output")
+    return (
+        result_a.output(parties[1], "my_output")
+        + result_b.output(parties[1], "my_output_b")
+        + result_c.output(parties[1], "my_output_c")
+        + result_d.output(parties[1], "my_output_d")
+    )

tests/nada-tests/src/matrix_multiplication_rational_multidim.py

@@ -1,21 +1,29 @@
 from nada_dsl import *
-
-# Step 0: Nada Algebra is imported with this line
 import nada_algebra as na
 
 
 def nada_main():
-    # Step 1: We use Nada Algebra wrapper to create "Party0", "Party1" and "Party2"
     parties = na.parties(3)
 
-    # Step 2: Party0 creates an array of dimension (2 x 1 x 1 x 2 x 2 x 2) with name "A"
     a = na.array([2, 1, 1, 2, 2, 2], parties[0], "A", na.SecretRational)
 
-    # Step 3: Party1 creates an array of dimension (2 x 1 x 1 x 2 x 2 x 2) with name "B"
     b = na.array([2, 1, 1, 2, 2, 2], parties[1], "B", na.SecretRational)
 
-    # Step 4: The result is of computing the dot product between the two which is another (2 x 1 x 1 x 2 x 2 x 2) matrix
-    result = a @ b
+    c = na.array((2,), parties[2], "C", na.SecretRational)
+
+    d = na.ones([2, 1, 1, 2, 2, 2], na.Rational)
+
+    result_a = a @ b
+
+    result_b = a @ c
+
+    result_c = a @ d
+
+    result_d = d @ a
 
-    # Step 5: We can use result.output() to produce the output for Party2 and variable name "my_output"
-    return result.output(parties[1], "my_output")
+    return (
+        result_a.output(parties[1], "my_output")
+        + result_b.output(parties[1], "my_output_b")
+        + result_c.output(parties[1], "my_output_c")
+        + result_d.output(parties[1], "my_output_d")
+    )

tests/nada-tests/tests/dot_product_rational.yaml

@@ -0,0 +1,26 @@
+---
+program: dot_product_rational
+inputs:
+  secrets:
+    A_0:
+      SecretInteger: "65536"
+    A_1:
+      SecretInteger: "131072"
+    A_2:
+      SecretInteger: "196608"
+    B_0:
+      SecretInteger: "65536"
+    B_1:
+      SecretInteger: "131072"
+    B_2:
+      SecretInteger: "196608"
+  public_variables: {}
+expected_outputs:
+  my_output_a_0:
+    SecretInteger: "917504"
+  my_output_b_0:
+    SecretInteger: "917504"
+  my_output_c_0:
+    SecretInteger: "393216"
+  my_output_d_0:
+    SecretInteger: "393216"

tests/nada-tests/tests/matrix_multiplication_rational.yaml

@@ -38,6 +38,12 @@ inputs:
       SecretInteger: "524288"
     B_2_2:
       SecretInteger: "589824"
+    C_0:
+      SecretInteger: "65536"
+    C_1:
+      SecretInteger: "0"
+    C_2:
+      SecretInteger: "0"
   public_variables: {}
 expected_outputs:
     my_output_0_0:
@@ -58,3 +64,33 @@ expected_outputs:
       SecretInteger: "8257536"
     my_output_2_2:
       SecretInteger: "9830400"
+    my_output_b_0:
+      SecretInteger: "65536"
+    my_output_b_1:
+      SecretInteger: "262144"
+    my_output_b_2:
+      SecretInteger: "458752"
+    my_output_c_0_0:
+      SecretInteger: "393216"
+    my_output_c_0_1:
+      SecretInteger: "393216"
+    my_output_c_0_2:
+      SecretInteger: "393216"
+    my_output_c_1_0:
+      SecretInteger: "983040"
+    my_output_c_1_1:
+      SecretInteger: "983040"
+    my_output_c_1_2:
+      SecretInteger: "983040"
+    my_output_c_2_0:
+      SecretInteger: "1572864"
+    my_output_c_2_1:
+      SecretInteger: "1572864"
+    my_output_c_2_2:
+      SecretInteger: "1572864"
+    my_output_d_0_0:
+      SecretInteger: "786432"
+    my_output_d_1_0:
+      SecretInteger: "786432"
+    my_output_d_2_0:
+      SecretInteger: "786432"