Add support for inplace encrypt

OpenMined · Dec 3, 2020 · 686e59d · 686e59d
1 parent fdb0188
commit 686e59d
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 18 deletions.
diff --git a/examples/experimental/FV HE Experimentation.ipynb b/examples/experimental/FV HE Experimentation.ipynb
@@ -17,13 +17,7 @@
     "from syft.frameworks.torch.he.fv.modulus import CoeffModulus\n",
     "from syft.frameworks.torch.he.fv.encryption_params import EncryptionParams\n",
     "from syft.frameworks.torch.he.fv.context import Context\n",
-    "from syft.frameworks.torch.he.fv.integer_encoder import IntegerEncoder\n",
     "from syft.frameworks.torch.he.fv.key_generator import KeyGenerator\n",
-    "from syft.frameworks.torch.he.fv.encryptor import Encryptor\n",
-    "from syft.frameworks.torch.he.fv.decryptor import Decryptor\n",
-    "from syft.frameworks.torch.he.fv.integer_encoder import IntegerEncoder\n",
-    "from syft.frameworks.torch.he.fv.modulus import SeqLevelType\n",
-    "from syft.frameworks.torch.he.fv.evaluator import Evaluator\n",
     "from syft.frameworks.torch.he.fv.context import Context\n",
     "from syft.frameworks.torch.he.fv.encryption_params import EncryptionParams\n",
     "context = Context(EncryptionParams(64, CoeffModulus().create(64, [30, 30]), 64))"

diff --git a/syft/frameworks/torch/tensors/interpreters/bfv.py b/syft/frameworks/torch/tensors/interpreters/bfv.py
@@ -37,23 +37,23 @@ def __init__(self, data=None, **kwargs):
         self.evaluator = None
 
     def encrypt(self, key, context):
-        self.context = context
         self.prepareEncryptor(key)
 
         inputs = self.child.flatten().tolist()
         new_child = [self.encryptor.encrypt(self.encoder.encode(x)) for x in inputs]
-        self.child = np.array(new_child).reshape(self.child.shape)
-        return self
-
-    def encrypt_(self, key, context):
         obj = BFVTensor()
         obj.context = context
+        obj.child = np.array(new_child).reshape(self.child.shape)
+        return obj
+
+    def encrypt_(self, key, context):
+        self.context = context
         self.prepareEncryptor(key)
 
         inputs = self.child.flatten().tolist()
         new_child = [self.encryptor.encrypt(self.encoder.encode(x)) for x in inputs]
-        obj.child = np.array(new_child).reshape(self.child.shape)
-        return obj
+        self.child = np.array(new_child).reshape(self.child.shape)
+        return self
 
     def decrypt(self, private_key):
         self.prepareDecryptor(private_key)

diff --git a/syft/frameworks/torch/tensors/interpreters/native.py b/syft/frameworks/torch/tensors/interpreters/native.py
@@ -1159,12 +1159,15 @@ def encrypt(self, protocol="mpc", inplace=False, **kwargs):
 
             x_encrypted = BFVTensor().on(self)
             if private_key is not None:
-                # Homomorphic Encryption using private key; if private key available prefer it.
-                x_encrypted.child.encrypt(private_key, context)
+                if not inplace:
+                    x_encrypted.child.encrypt(private_key, context)
+                else:
+                    x_encrypted.child.encrypt_(private_key, context)
             elif public_key is not None:
-                x_encrypted.child.encrypt(
-                    public_key, context
-                )  # Homomorphic Encryption using public key
+                if not inplace:
+                    x_encrypted.child.encrypt(public_key, context)
+                else:
+                    x_encrypted.child.encrypt_(public_key, context)
             else:
                 raise RuntimeError("Secret key or Public key should be provided for encryption")
             return x_encrypted