Repository: OpenMined/PySyft

Issue:
Implement Relinearization keys for FV scheme
## Feature Description
**Relinearization keys** need to be implemented for the FV scheme.
These keys are useful for reducing the number of polynomials produced after the multiplication operations of two ciphertexts.


Discussion:
This issue has been marked stale because it has been open 30 days with no activity. Leave a comment or remove the `stale` label to unmark it. Otherwise, this will be closed in 7 days.
This issue has been marked stale because it has been open 30 days with no activity. Leave a comment or remove the `stale` label to unmark it. Otherwise, this will be closed in 7 days.

Resolving patch:
diff --git a/syft/frameworks/torch/he/fv/evaluator.py b/syft/frameworks/torch/he/fv/evaluator.py
--- a/syft/frameworks/torch/he/fv/evaluator.py
+++ b/syft/frameworks/torch/he/fv/evaluator.py
@@ -11,6 +11,7 @@
 from syft.frameworks.torch.he.fv.util.operations import multiply_sub_plain_with_delta
 from syft.frameworks.torch.he.fv.ciphertext import CipherText
 from syft.frameworks.torch.he.fv.plaintext import PlainText
+from syft.frameworks.torch.he.fv.relin_keys import RelinKey
 
 
 class ParamTypes(Enum):
@@ -381,3 +382,103 @@ def _mul_plain_plain(self, pt1, pt2):
         """
         pt1, pt2 = pt1.data, pt2.data
         return PlainText(poly_mul_mod(pt1, pt2, self.plain_modulus, self.poly_modulus))
+
+    def relin(self, ct, key):
+        """Relinearize the provided ciphertext and decrease its size by one.
+        (cannot apply on size 2 ciphetext)
+
+        Args:
+            ct (CipherText): ciphertext of size 3 to relinearize.
+            key (Relin Key): relinearization key generated with keygenerator.
+
+        Returns:
+            Ciphertext of Size 2 with same encrypted value.
+        """
+        if len(ct.data) == 2:
+            raise Warning("Ciphertext of size 2 does not need to relinearize.")
+        if len(ct.data) > 3:
+            raise Warning(f"Ciphertext of size {len(ct.data)} cannot be relinearized.")
+
+        return self._switch_key_inplace(copy.deepcopy(ct), key)
+
+    def _switch_key_inplace(self, ct, key):
+
+        if not isinstance(key, RelinKey):
+            raise RuntimeError("Relinearization key is invalid")
+
+        param_id = ct.param_id
+        ct = ct.data
+        key_vector = key.data
+        context_data = self.context.context_data_map[param_id]
+        key_context = self.context.context_data_map[self.context.key_param_id]
+
+        coeff_modulus = context_data.param.coeff_modulus
+        decomp_mod_count = len(coeff_modulus)
+        key_mod = key_context.param.coeff_modulus
+        key_mod_count = len(key_mod)
+        rns_mod_count = decomp_mod_count + 1
+
+        target = ct[-1]  # Last component of ciphertext
+
+        modswitch_factors = key_context.rns_tool.inv_q_last_mod_q
+
+        for i in range(decomp_mod_count):
+
+            local_small_poly_0 = copy.deepcopy(target[i])
+
+            temp_poly = [[[0] for x in range(rns_mod_count)], [[0] for x in range(rns_mod_count)]]
+
+            for j in range(rns_mod_count):
+                index = key_mod_count - 1 if j == decomp_mod_count else j
+
+                if key_mod[i] <= key_mod[index]:
+                    local_small_poly_1 = copy.deepcopy(local_small_poly_0)
+                else:
+                    local_small_poly_1 = [x % key_mod[index] for x in local_small_poly_0]
+
+                for k in range(2):
+                    local_small_poly_2 = poly_mul_mod(
+                        local_small_poly_1,
+                        key_vector[i][k][index],
+                        key_mod[index],
+                        self.poly_modulus,
+                    )
+                    temp_poly[k][j] = poly_add_mod(
+                        local_small_poly_2, temp_poly[k][j], key_mod[index], self.poly_modulus
+                    )
+
+        # Results are now stored in temp_poly[k]
+        # Modulus switching should be performed
+        for k in range(2):
+            temp_poly_ptr = temp_poly[k][decomp_mod_count]
+            temp_last_poly_ptr = temp_poly[k][decomp_mod_count]
+
+            temp_poly_ptr = [x % key_mod[-1] for x in temp_poly_ptr]
+
+            # Add (p-1)/2 to change from flooring to rounding.
+            half = key_mod[-1] >> 1
+            temp_last_poly_ptr = [(x + half) % key_mod[-1] for x in temp_last_poly_ptr]
+
+            encrypted_ptr = ct[k]
+            for j in range(decomp_mod_count):
+                temp_poly_ptr = temp_poly[k][j]
+
+                temp_poly_ptr = [x % key_mod[j] for x in temp_poly_ptr]
+                local_small_poly = [x % key_mod[j] for x in temp_last_poly_ptr]
+                half_mod = half % key_mod[j]
+
+                local_small_poly = [(x - half_mod) % key_mod[j] for x in local_small_poly]
+
+                # ((ct mod qi) - (ct mod qk)) mod qi
+                temp_poly_ptr = poly_sub_mod(
+                    temp_poly_ptr, local_small_poly, key_mod[j], self.poly_modulus
+                )
+
+                # qk^(-1) * ((ct mod qi) - (ct mod qk)) mod qi
+                temp_poly_ptr = [(x * modswitch_factors[j]) % key_mod[j] for x in temp_poly_ptr]
+
+                encrypted_ptr[j] = poly_add_mod(
+                    temp_poly_ptr, encrypted_ptr[j], key_mod[j], self.poly_modulus
+                )
+
+        return CipherText(ct[0:2], param_id)
diff --git a/syft/frameworks/torch/he/fv/key_generator.py b/syft/frameworks/torch/he/fv/key_generator.py
--- a/syft/frameworks/torch/he/fv/key_generator.py
+++ b/syft/frameworks/torch/he/fv/key_generator.py
@@ -3,6 +3,7 @@
 from syft.frameworks.torch.he.fv.util.rlwe import encrypt_symmetric
 from syft.frameworks.torch.he.fv.secret_key import SecretKey
 from syft.frameworks.torch.he.fv.public_key import PublicKey
+from syft.frameworks.torch.he.fv.relin_keys import RelinKeys
 
 
 class KeyGenerator:
@@ -16,11 +17,12 @@ class KeyGenerator:
 
     def __init__(self, context):
         if not isinstance(context, Context):
-            raise ValueError("invalid context")
+            raise RuntimeError("invalid context")
 
         self.public_key = None
         self.secret_key = None
         self.context = context
+        self.relin_key_generator = None
 
     def keygen(self):
         """Generate the secret key and public key.
@@ -44,3 +46,18 @@ def _generate_pk(self):
             self.context, self.context.key_param_id, self.secret_key.data
         )
         self.public_key = PublicKey(public_key.data)
+
+    def get_relin_keys(self):
+        """Generate a relinearization key.
+
+        Returns:
+            A relinearization key.
+        """
+        if self.relin_key_generator is None:
+            if self.secret_key is None:
+                raise RuntimeError("cannot generate relinearization key for unspecified secret key")
+
+            self.relin_key_generator = RelinKeys(self.context, self.secret_key)
+
+        # generate keys.
+        return self.relin_key_generator._generate_relin_key()
diff --git a/syft/frameworks/torch/he/fv/relin_keys.py b/syft/frameworks/torch/he/fv/relin_keys.py
new file mode 100644
--- /dev/null
+++ b/syft/frameworks/torch/he/fv/relin_keys.py
@@ -0,0 +1,45 @@
+from syft.frameworks.torch.he.fv.util.operations import poly_mul_mod
+from syft.frameworks.torch.he.fv.util.rlwe import encrypt_symmetric
+from syft.frameworks.torch.he.fv.util.operations import poly_add_mod
+
+
+class RelinKey:
+    def __init__(self, data):
+        self.data = data
+
+
+class RelinKeys:
+    def __init__(self, context, sk):
+        self._context = context
+        key_param = context.context_data_map[context.key_param_id].param
+        self._coeff_modulus = key_param.coeff_modulus
+        self._poly_modulus = key_param.poly_modulus
+        self._secret_key = sk.data
+        self._secret_key_power_2 = self._get_sk_power_2(sk.data)
+
+    def _generate_relin_key(self):
+        return RelinKey(self._generate_one_kswitch_key(self._secret_key_power_2))
+
+    def _generate_one_kswitch_key(self, sk_power_2):
+        decomp_mod_count = len(self._coeff_modulus) - 1
+        result = [0] * (decomp_mod_count)
+        for i in range(decomp_mod_count):
+            result[i] = encrypt_symmetric(
+                self._context, self._context.key_param_id, self._secret_key
+            ).data
+            factor = self._coeff_modulus[-1] % self._coeff_modulus[i]
+
+            temp = [(x * factor) for x in sk_power_2[i]]
+
+            result[i][0][i] = poly_add_mod(
+                result[i][0][i], temp, self._coeff_modulus[i], self._poly_modulus
+            )
+        return result
+
+    def _get_sk_power_2(self, sk):
+        sk_power_2 = []
+        for i in range(len(self._coeff_modulus)):
+            sk_power_2.append(
+                poly_mul_mod(sk[i], sk[i], self._coeff_modulus[i], self._poly_modulus)
+            )
+        return sk_power_2


Test patch:
diff --git a/test/torch/tensors/test_fv.py b/test/torch/tensors/test_fv.py
--- a/test/torch/tensors/test_fv.py
+++ b/test/torch/tensors/test_fv.py
@@ -676,3 +676,47 @@ def test_rns_tool_fastbconv_sk(poly_len, coeff_mod, plain_mod, input, output):
     rns_tool = RNSTool(enc_param)
     result = rns_tool.fastbconv_sk(input)
     assert result == output
+
+
+@pytest.mark.parametrize(
+    "val1, val2", [(0, 0), (1, 1), (-1, 1), (100, -1), (1000, 1), (-1000, -1), (-99, 0)],
+)
+def test_fv_relin(val1, val2):
+    ctx = Context(EncryptionParams(128, CoeffModulus().create(128, [40, 40]), 64))
+    keygenerator = KeyGenerator(ctx)
+    keys = keygenerator.keygen()
+    relin_key = keygenerator.get_relin_keys()
+    encoder = IntegerEncoder(ctx)
+    encryptor = Encryptor(ctx, keys[1])  # keys[1] = public_key
+    decryptor = Decryptor(ctx, keys[0])  # keys[0] = secret_key
+    evaluator = Evaluator(ctx)
+
+    op1 = encryptor.encrypt(encoder.encode(val1))
+    op2 = encryptor.encrypt(encoder.encode(val2))
+    temp_prod = evaluator.mul(op1, op2)
+    relin_prod = evaluator.relin(temp_prod, relin_key)
+    assert len(temp_prod.data) - 1 == len(relin_prod.data)
+    assert val1 * val2 == encoder.decode(decryptor.decrypt(relin_prod))
+
+
+@pytest.mark.parametrize(
+    "val1, val2", [(-1, 1)],
+)
+def test_fv_relin_exceptions(val1, val2):
+    ctx = Context(EncryptionParams(128, CoeffModulus().create(128, [40, 40]), 64))
+    keygenerator = KeyGenerator(ctx)
+    keys = keygenerator.keygen()
+    relin_key = keygenerator.get_relin_keys()
+    encoder = IntegerEncoder(ctx)
+    encryptor = Encryptor(ctx, keys[1])  # keys[1] = public_key
+    evaluator = Evaluator(ctx)
+
+    op1 = encryptor.encrypt(encoder.encode(val1))
+    op2 = encryptor.encrypt(encoder.encode(val2))
+    temp_prod = evaluator.mul(op1, op2)
+
+    with pytest.raises(Warning):
+        evaluator.relin(op1, relin_key)  # Ciphertext size 2
+
+    with pytest.raises(Exception):
+        evaluator.relin(evaluator.mul(temp_prod, val1), relin_key)  # Ciphertext size 4