fix bug from assuming that inputs were zero-dimensional

cicada-mpc · Jan 31, 2024 · 2ac4ec1 · 2ac4ec1
1 parent 9f54f80
commit 2ac4ec1
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 8 deletions.
diff --git a/cicada/additive.py b/cicada/additive.py
@@ -332,16 +332,25 @@ def divide(self, lhs, rhs, *, encoding=None, rmask=None, mask1=None, rem1=None,
         # Private-private division.
         if isinstance(lhs, AdditiveArrayShare) and isinstance(rhs, AdditiveArrayShare):
             zshare = self.share(src=0, secret=numpy.zeros_like(rhs.storage), shape=rhs.storage.shape)
-            if self.reveal(self.equal(rhs, zshare), encoding=Boolean()):
-                raise ZeroDivisionError()
+            check = self.reveal(self.equal(rhs, zshare), encoding=Boolean())
+            if not len(check.shape):
+                if check:
+                    raise ZeroDivisionError()
+            else:
+                if check.any():   
+                    raise ZeroDivisionError()
             oops = True
             while oops:
                 if rmask is None:
                     _, rmask = self.random_bitwise_secret(bits=encoding.precision, shape=rhs.storage.shape)
-                if self.reveal(self.equal(rmask, zshare), encoding=Boolean()):
-                    rmask = None
                 else:
                     oops = False
+                if not len(rhs.storage.shape):
+                    if self.reveal(self.equal(rmask, zshare), encoding=Boolean()):
+                        rmask = None
+                else:
+                    if self.reveal(self.equal(rmask, zshare), encoding=Boolean()).any():
+                        rmask = None
             rhsmasked = self.field_multiply(rmask, rhs)
             if mask1 != None and rem1 != None:
                 rhsmasked = self.right_shift(rhsmasked, bits=encoding.precision, trunc_mask=mask1, rem_mask=rem1)

diff --git a/cicada/shamir.py b/cicada/shamir.py
@@ -798,16 +798,25 @@ def divide(self, lhs, rhs, *, encoding=None, rmask=None, mask1=None, rem1=None,
         # Private-private division.
         if isinstance(lhs, ShamirArrayShare) and isinstance(rhs, ShamirArrayShare):
             zshare = self.share(src=0, secret=numpy.zeros_like(rhs.storage), shape=rhs.storage.shape)
-            if self.reveal(self.equal(rhs, zshare), encoding=Boolean()):
-                raise ZeroDivisionError()
+            check = self.reveal(self.equal(rhs, zshare), encoding=Boolean())
+            if not len(check.shape):
+                if check:
+                    raise ZeroDivisionError()
+            else:
+                if check.any():   
+                    raise ZeroDivisionError()
             oops = True
             while oops:
                 if rmask is None:
                     _, rmask = self.random_bitwise_secret(bits=encoding.precision, shape=rhs.storage.shape)
-                if self.reveal(self.equal(rmask, zshare), encoding=Boolean()):
-                    rmask = None
                 else:
                     oops = False
+                if not len(rhs.storage.shape):
+                    if self.reveal(self.equal(rmask, zshare), encoding=Boolean()):
+                        rmask = None
+                else:
+                    if self.reveal(self.equal(rmask, zshare), encoding=Boolean()).any():
+                        rmask = None
             rhsmasked = self.field_multiply(rmask, rhs)
             if mask1 != None and rem1 != None:
                 rhsmasked = self.right_shift(rhsmasked, bits=encoding.precision, trunc_mask=mask1, rem_mask=rem1)