Various small bug fixes (#12)

Various small bug fixes

Author: mathias-nillion

nada_algebra/array.py

@@ -163,6 +163,24 @@ def __mul__(self, other: _NadaOperand) -> "NadaArray":
             return NadaArray(self.inner * Integer(other))
         return NadaArray(self.inner * other)
 
+    def __pow__(self, other: int) -> "NadaArray":
+        """Raises NadaArray to a power.
+
+        Args:
+            other (int): Power value.
+
+        Returns:
+            NadaArray: Result NadaArray.
+        """
+        if not isinstance(other, int):
+            raise TypeError(
+                "Cannot raise `NadaArray` to power of type `%s`" % type(other).__name__
+            )
+        result = self.copy()
+        for _ in range(other - 1):
+            result = result * result
+        return result
+
     def __truediv__(self, other: _NadaOperand) -> "NadaArray":
         """
         Perform element-wise division with broadcasting.
@@ -189,8 +207,11 @@ def __matmul__(self, other: "NadaArray") -> "NadaArray":
         Returns:
             NadaArray: A new NadaArray representing the result of matrix multiplication.
         """
-        if isinstance(other, NadaArray):
-            return NadaArray(self.inner @ other.inner)
+        return NadaArray(self.inner @ other.inner)
+
+    @property
+    def ndim(self) -> int:
+        return len(self.shape)
 
     def dot(self, other: "NadaArray") -> "NadaArray":
         """
@@ -295,10 +316,17 @@ def output_array(array: np.ndarray, party: Party, prefix: str) -> list:
             ),
         ):
             return [Output(array, f"{prefix}_0", party)]
+        elif isinstance(array, (Rational, SecretRational)):
+            return [Output(array.value, f"{prefix}_0", party)]
 
         if len(array.shape) == 1:
             return [
-                Output(array[i], f"{prefix}_{i}", party) for i in range(array.shape[0])
+                (
+                    Output(array[i].value, f"{prefix}_{i}", party)
+                    if isinstance(array[i], (Rational, SecretRational))
+                    else Output(array[i], f"{prefix}_{i}", party)
+                )
+                for i in range(array.shape[0])
             ]
         return [
             v

nada_algebra/types.py

@@ -90,7 +90,7 @@ def from_number(
         Returns:
             Rational: Instantiated wrapper around number.
         """
-        if (value is None) or (value == np.nan):
+        if value is None:
             raise ValueError("Cannot convert `%s` to Rational." % value)
 
         value = value.item() if isinstance(value, np.floating) else value
@@ -308,13 +308,13 @@ def rescale(value: NadaType, scale: UnsignedInteger, direction: str) -> NadaType
         try:
             return value << scale
         except:
-            return value * (1 << scale)
+            return value * Integer(1 << scale)
     elif direction == "down":
         # TODO: remove try block when rshift implemented for every NadaType
         try:
             return value >> scale
         except:
-            return value / (1 << scale)
+            return value / Integer(1 << scale)
 
     raise ValueError(
         'Invalid scaling direction `%s`. Expected "up" or "down"' % direction