feat: added rational-specific matrix multiplication (#25)

* feat: added rational specific matrix multiplication

* feat: added support for multidimensional rational matrices for multiplications

* fix: float_from_rational allow input log_scale

* fix: various small fixes

Author: jcabrero

nada_algebra/array.py

@@ -243,9 +243,75 @@ def matmul(self, other: "NadaArray") -> "NadaArray":
             NadaArray: A new NadaArray representing the result of matrix multiplication.
         """
         if isinstance(other, NadaArray):
+            if self.is_rational or other.is_rational:
+                return self.rational_matmul(other)
             return NadaArray(np.array(self.inner @ other.inner))
         return NadaArray(np.array(self.inner @ other))
 
+    def rational_matmul(self, other: "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.
+        """
+        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}"
+            )
+
+        # 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.

nada_algebra/client.py

@@ -135,3 +135,19 @@ def secret_rational(value: Union[float, int]) -> SecretInteger:
         int: The integer representation of the input value.
     """
     return SecretInteger(__rational(value))
+
+
+def float_from_rational(value: int, log_scale: int = None) -> float:
+    """
+    Returns the float representation of the given rational value.
+
+    Args:
+        value (int): The output Rational value to convert.
+        log_scale (int, optional): The log scale to use for conversion. Defaults to None.
+
+    Returns:
+        float: The float representation of the input value.
+    """
+    if log_scale is not None:
+        log_scale = get_log_scale()
+    return value / (1 << log_scale)

tests/nada-tests/nada-project.toml

@@ -145,3 +145,11 @@ prime_size = 128
 [[programs]]
 path = "src/array_statistics.py"
 prime_size = 128
+
+[[programs]]
+path = "src/matrix_multiplication_rational.py"
+prime_size = 128
+
+[[programs]]
+path = "src/matrix_multiplication_rational_multidim.py"
+prime_size = 128

tests/nada-tests/src/matrix_multiplication_rational.py

@@ -0,0 +1,21 @@
+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
+
+    # 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")

tests/nada-tests/src/matrix_multiplication_rational_multidim.py

@@ -0,0 +1,21 @@
+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
+
+    # 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")

tests/nada-tests/tests/matrix_multiplication_rational.yaml

@@ -0,0 +1,60 @@
+---
+program: matrix_multiplication_rational
+inputs:
+  secrets:
+    A_0_0:
+      SecretInteger: "65536"
+    A_0_1:
+      SecretInteger: "131072"
+    A_0_2:
+      SecretInteger: "196608"
+    A_1_0:
+      SecretInteger: "262144"
+    A_1_1:
+      SecretInteger: "327680"
+    A_1_2:
+      SecretInteger: "393216"
+    A_2_0:
+      SecretInteger: "458752"
+    A_2_1:
+      SecretInteger: "524288"
+    A_2_2:
+      SecretInteger: "589824"
+    B_0_0:
+      SecretInteger: "65536"
+    B_0_1:
+      SecretInteger: "131072"
+    B_0_2:
+      SecretInteger: "196608"
+    B_1_0:
+      SecretInteger: "262144"
+    B_1_1:
+      SecretInteger: "327680"
+    B_1_2:
+      SecretInteger: "393216"
+    B_2_0:
+      SecretInteger: "458752"
+    B_2_1:
+      SecretInteger: "524288"
+    B_2_2:
+      SecretInteger: "589824"
+  public_variables: {}
+expected_outputs:
+    my_output_0_0:
+      SecretInteger: "1966080"
+    my_output_0_1:
+      SecretInteger: "2359296"
+    my_output_0_2:
+      SecretInteger: "2752512"
+    my_output_1_0:
+      SecretInteger: "4325376"
+    my_output_1_1:
+      SecretInteger: "5308416"
+    my_output_1_2:
+      SecretInteger: "6291456"
+    my_output_2_0:
+      SecretInteger: "6684672"
+    my_output_2_1:
+      SecretInteger: "8257536"
+    my_output_2_2:
+      SecretInteger: "9830400"

tests/nada-tests/tests/matrix_multiplication_rational_multidim.yaml

@@ -0,0 +1,102 @@
+---
+program: matrix_multiplication_rational_multidim
+inputs:
+  secrets:
+    A_0_0_0_0_0_0:
+      SecretInteger: "65536"
+    A_0_0_0_0_0_1:
+      SecretInteger: "131072"
+    A_0_0_0_0_1_0:
+      SecretInteger: "196608"
+    A_0_0_0_0_1_1:
+      SecretInteger: "262144"
+    A_0_0_0_1_0_0:
+      SecretInteger: "65536"
+    A_0_0_0_1_0_1:
+      SecretInteger: "131072"
+    A_0_0_0_1_1_0:
+      SecretInteger: "196608"
+    A_0_0_0_1_1_1:
+      SecretInteger: "262144"
+    A_1_0_0_0_0_0:
+      SecretInteger: "65536"
+    A_1_0_0_0_0_1:
+      SecretInteger: "131072"
+    A_1_0_0_0_1_0:
+      SecretInteger: "196608"
+    A_1_0_0_0_1_1:
+      SecretInteger: "262144"
+    A_1_0_0_1_0_0:
+      SecretInteger: "65536"
+    A_1_0_0_1_0_1:
+      SecretInteger: "131072"
+    A_1_0_0_1_1_0:
+      SecretInteger: "196608"
+    A_1_0_0_1_1_1:
+      SecretInteger: "262144"
+    B_0_0_0_0_0_0:
+      SecretInteger: "65536"
+    B_0_0_0_0_0_1:
+      SecretInteger: "131072"
+    B_0_0_0_0_1_0:
+      SecretInteger: "196608"
+    B_0_0_0_0_1_1:
+      SecretInteger: "262144"
+    B_0_0_0_1_0_0:
+      SecretInteger: "65536"
+    B_0_0_0_1_0_1:
+      SecretInteger: "131072"
+    B_0_0_0_1_1_0:
+      SecretInteger: "196608"
+    B_0_0_0_1_1_1:
+      SecretInteger: "262144"
+    B_1_0_0_0_0_0:
+      SecretInteger: "65536"
+    B_1_0_0_0_0_1:
+      SecretInteger: "131072"
+    B_1_0_0_0_1_0:
+      SecretInteger: "196608"
+    B_1_0_0_0_1_1:
+      SecretInteger: "262144"
+    B_1_0_0_1_0_0:
+      SecretInteger: "65536"
+    B_1_0_0_1_0_1:
+      SecretInteger: "131072"
+    B_1_0_0_1_1_0:
+      SecretInteger: "196608"
+    B_1_0_0_1_1_1:
+      SecretInteger: "262144"
+  public_variables: {}
+expected_outputs:
+  my_output_0_0_0_0_0_0:
+    SecretInteger: "458752"
+  my_output_0_0_0_0_0_1:
+    SecretInteger: "655360"
+  my_output_0_0_0_0_1_0:
+    SecretInteger: "983040"
+  my_output_0_0_0_0_1_1:
+    SecretInteger: "1441792"
+  my_output_0_0_0_1_0_0:
+    SecretInteger: "458752"
+  my_output_0_0_0_1_0_1:
+    SecretInteger: "655360"
+  my_output_0_0_0_1_1_0:
+    SecretInteger: "983040"
+  my_output_0_0_0_1_1_1:
+    SecretInteger: "1441792"
+  my_output_1_0_0_0_0_0:
+    SecretInteger: "458752"
+  my_output_1_0_0_0_0_1:
+    SecretInteger: "655360"
+  my_output_1_0_0_0_1_0:
+    SecretInteger: "983040"
+  my_output_1_0_0_0_1_1:
+    SecretInteger: "1441792"
+  my_output_1_0_0_1_0_0:
+    SecretInteger: "458752"
+  my_output_1_0_0_1_0_1:
+    SecretInteger: "655360"
+  my_output_1_0_0_1_1_0:
+    SecretInteger: "983040"
+  my_output_1_0_0_1_1_1:
+    SecretInteger: "1441792"

tests/test_all.py

@@ -37,6 +37,8 @@
     "array_attributes",
     "functional_operations",
     "array_statistics",
+    "matrix_multiplication_rational",
+    "matrix_multiplication_rational_multidim",
     # Not supported yet
     # "unsigned_matrix_inverse",
     # "private_inverse"