390 lines
13 KiB
Python
390 lines
13 KiB
Python
# built-in dependencies
|
|
from typing import Union, List
|
|
import multiprocessing
|
|
from contextlib import closing
|
|
|
|
# 3rd party dependencies
|
|
from tqdm import tqdm
|
|
|
|
# project dependencies
|
|
from lightphe.models.Homomorphic import Homomorphic
|
|
from lightphe.commons import phe_utils
|
|
from lightphe.models.Ciphertext import Ciphertext
|
|
from lightphe.models.Algorithm import Algorithm
|
|
from lightphe.commons.logger import Logger
|
|
|
|
logger = Logger(module="lightphe/models/Tensor.py")
|
|
|
|
|
|
# pylint: disable=too-few-public-methods, no-else-return
|
|
class Fraction:
|
|
"""
|
|
Class to store fractional values
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dividend: Union[int, tuple, list],
|
|
abs_dividend: Union[int, tuple, list],
|
|
divisor: Union[int, tuple, list],
|
|
sign: int = 1,
|
|
):
|
|
self.dividend = dividend
|
|
self.divisor = divisor
|
|
self.sign = sign
|
|
self.abs_dividend = abs_dividend
|
|
|
|
def __str__(self):
|
|
"""
|
|
Print Fraction Class Object
|
|
"""
|
|
sign = "-" if self.sign == -1 else "+"
|
|
return f"Fraction({sign}{self.abs_dividend} / {self.divisor})"
|
|
|
|
def __repr__(self):
|
|
"""
|
|
Print Fraction Class Object
|
|
"""
|
|
return self.__str__()
|
|
|
|
|
|
class EncryptedTensor:
|
|
"""
|
|
Class to store encrypted tensor objects
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
fractions: List[Fraction],
|
|
cs: Homomorphic,
|
|
precision: int = 5,
|
|
):
|
|
"""
|
|
Initialization method
|
|
Args:
|
|
fractions (list): list of fractions storing individual encrypted tensor items
|
|
cs (cryptosystem): built cryptosystem
|
|
precision (int): precision of the tensor
|
|
"""
|
|
self.fractions = fractions
|
|
self.cs = cs
|
|
self.precision = precision
|
|
|
|
def __str__(self):
|
|
"""
|
|
Print encrypted tensor object
|
|
"""
|
|
results = []
|
|
for i in self.fractions:
|
|
results.append(f"{i}")
|
|
return ", ".join(results)
|
|
|
|
def __repr__(self):
|
|
"""
|
|
Print encrypted tensor object
|
|
"""
|
|
return self.__str__()
|
|
|
|
def __matmul__(self, other: Union["EncryptedTensor", list]):
|
|
"""
|
|
Perform dot product of two encrypted tensors
|
|
"""
|
|
if not (
|
|
(isinstance(other, EncryptedTensor) and isinstance(self, list))
|
|
or (isinstance(other, list) and isinstance(self, EncryptedTensor))
|
|
):
|
|
raise ValueError(
|
|
"Dot product can be run for EncryptedTensor and List of float / int"
|
|
)
|
|
|
|
encrypted_tensor = self.__mul__(other=other)
|
|
|
|
if len(encrypted_tensor.fractions) == 0:
|
|
raise ValueError("Dot product cannot be calculated for empty tensor")
|
|
|
|
divisor = cast_ciphertext(
|
|
cs=self.cs, value=encrypted_tensor.fractions[0].divisor
|
|
)
|
|
|
|
if len(encrypted_tensor.fractions) > 10000:
|
|
# parallelize the sum operation
|
|
num_workers = min(
|
|
len(encrypted_tensor.fractions), multiprocessing.cpu_count()
|
|
)
|
|
|
|
chunks = chunkify(encrypted_tensor.fractions, num_workers)
|
|
|
|
with closing(multiprocessing.Pool(num_workers)) as pool:
|
|
funclist = []
|
|
|
|
for chunk in chunks:
|
|
f = pool.apply_async(sum_fractions_chunk, (chunk, self.cs))
|
|
funclist.append(f)
|
|
|
|
partial_sums = []
|
|
for f in tqdm(funclist, desc="Summing up fractions", disable=True):
|
|
result = f.get(timeout=10)
|
|
partial_sums.append(result)
|
|
|
|
# map reduce
|
|
total_sum = partial_sums[0]
|
|
for partial in partial_sums[1:]:
|
|
total_sum += partial
|
|
|
|
fraction = Fraction(
|
|
dividend=total_sum.value,
|
|
abs_dividend=total_sum.value,
|
|
divisor=divisor.value,
|
|
sign=1,
|
|
)
|
|
else:
|
|
# serial implementation
|
|
sum_dividend = cast_ciphertext(
|
|
cs=self.cs, value=encrypted_tensor.fractions[0].abs_dividend
|
|
)
|
|
divisor = cast_ciphertext(
|
|
cs=self.cs, value=encrypted_tensor.fractions[0].divisor
|
|
)
|
|
|
|
if len(encrypted_tensor.fractions) > 1:
|
|
for fraction in encrypted_tensor.fractions[1:]:
|
|
sum_dividend += cast_ciphertext(
|
|
value=fraction.abs_dividend, cs=self.cs
|
|
)
|
|
|
|
fraction = Fraction(
|
|
dividend=sum_dividend.value,
|
|
abs_dividend=sum_dividend.value,
|
|
divisor=divisor.value,
|
|
sign=1,
|
|
)
|
|
|
|
return EncryptedTensor(fractions=[fraction], cs=self.cs)
|
|
|
|
def __mul__(
|
|
self, other: Union["EncryptedTensor", int, float, list]
|
|
) -> "EncryptedTensor":
|
|
"""
|
|
Perform homomorphic element-wise multipliction on tensors
|
|
or multiplication of an encrypted tensor with a constant
|
|
Args:
|
|
other: encrypted tensor or constant
|
|
Returns:
|
|
encrypted tensor
|
|
"""
|
|
if isinstance(other, EncryptedTensor):
|
|
if isinstance(other, EncryptedTensor) and len(self.fractions) != len(
|
|
other.fractions
|
|
):
|
|
raise ValueError(
|
|
"Tensor sizes must be equal in homomorphic multiplication"
|
|
)
|
|
|
|
fractions = []
|
|
for i, alpha_tensor in enumerate(self.fractions):
|
|
beta_tensor = other.fractions[i]
|
|
|
|
current_dividend = self.cs.multiply(
|
|
ciphertext1=alpha_tensor.dividend, ciphertext2=beta_tensor.dividend
|
|
)
|
|
|
|
current_abs_dividend = self.cs.multiply(
|
|
ciphertext1=alpha_tensor.abs_dividend,
|
|
ciphertext2=beta_tensor.abs_dividend,
|
|
)
|
|
|
|
current_divisor = self.cs.multiply(
|
|
ciphertext1=alpha_tensor.divisor, ciphertext2=beta_tensor.divisor
|
|
)
|
|
|
|
fraction = Fraction(
|
|
dividend=current_dividend,
|
|
abs_dividend=current_abs_dividend,
|
|
divisor=current_divisor,
|
|
sign=alpha_tensor.sign * beta_tensor.sign,
|
|
)
|
|
|
|
fractions.append(fraction)
|
|
|
|
return EncryptedTensor(fractions=fractions, cs=self.cs)
|
|
|
|
elif isinstance(other, list):
|
|
# perform element-wise multiplication of encrypted tensor with plain tensor
|
|
if len(self.fractions) != len(other):
|
|
raise ValueError(
|
|
"Tensor sizes must be equal in homomorphic multiplication"
|
|
)
|
|
|
|
if any(i < 0 for i in other) or any(
|
|
fraction.sign < 0 for fraction in self.fractions
|
|
):
|
|
raise ValueError(
|
|
"all items in the plain and encrypted tensor must be positive"
|
|
" to perform element wise multiplication"
|
|
)
|
|
|
|
dividends = []
|
|
divisor = None
|
|
for alpha, beta in zip(self.fractions, other):
|
|
c_abs_dividend, c_divisor = phe_utils.fractionize(
|
|
value=(
|
|
abs(beta) % self.cs.plaintext_modulo
|
|
if abs(beta) > self.cs.plaintext_modulo
|
|
else abs(beta)
|
|
),
|
|
modulo=self.cs.plaintext_modulo,
|
|
precision=self.precision,
|
|
)
|
|
|
|
dividend = (
|
|
cast_ciphertext(cs=self.cs, value=alpha.abs_dividend)
|
|
* c_abs_dividend
|
|
)
|
|
|
|
if divisor is None:
|
|
divisor = (
|
|
cast_ciphertext(cs=self.cs, value=alpha.divisor) * c_divisor
|
|
)
|
|
|
|
# dividends.append(dividend)
|
|
dividends.append(
|
|
Fraction(
|
|
dividend=dividend.value,
|
|
abs_dividend=dividend.value,
|
|
divisor=divisor.value,
|
|
sign=1,
|
|
)
|
|
)
|
|
|
|
return EncryptedTensor(fractions=dividends, cs=self.cs)
|
|
|
|
elif isinstance(other, (int, float)):
|
|
constant_sign = 1 if other >= 0 else -1
|
|
other = abs(other)
|
|
if isinstance(other, float):
|
|
other = phe_utils.normalize_input(
|
|
value=other, modulo=self.cs.plaintext_modulo
|
|
)
|
|
|
|
fractions = []
|
|
for alpha_tensor in self.fractions:
|
|
dividend = self.cs.multiply_by_constant(
|
|
ciphertext=alpha_tensor.dividend, constant=other
|
|
)
|
|
abs_dividend = self.cs.multiply_by_constant(
|
|
ciphertext=alpha_tensor.abs_dividend, constant=other
|
|
)
|
|
# notice that divisor is alpha tensor's divisor instead of addition
|
|
fraction = Fraction(
|
|
dividend=dividend,
|
|
abs_dividend=abs_dividend,
|
|
divisor=alpha_tensor.divisor,
|
|
sign=constant_sign * alpha_tensor.sign,
|
|
)
|
|
fractions.append(fraction)
|
|
return EncryptedTensor(fractions=fractions, cs=self.cs)
|
|
else:
|
|
raise ValueError(
|
|
"Encrypted tensor can be multiplied by an encrypted tensor or constant"
|
|
)
|
|
|
|
def __rmul__(
|
|
self, multiplier: Union[int, float, "EncryptedTensor"]
|
|
) -> "EncryptedTensor":
|
|
"""
|
|
Perform multiplication of encrypted tensor with a constant or plain tensor (element-wise)
|
|
Args:
|
|
multiplier: scalar value
|
|
Returns:
|
|
encrypted tensor
|
|
"""
|
|
return self.__mul__(other=multiplier)
|
|
|
|
def __add__(self, other: "EncryptedTensor") -> "EncryptedTensor":
|
|
"""
|
|
Perform homomorphic addition
|
|
Args:
|
|
other: encrypted tensor
|
|
Returns:
|
|
encrypted tensor
|
|
"""
|
|
if len(self.fractions) != len(other.fractions):
|
|
raise ValueError("Fraction sizes must be equal")
|
|
|
|
current_tensors = []
|
|
for i, alpha_tensor in enumerate(self.fractions):
|
|
beta_tensor = other.fractions[i]
|
|
|
|
current_dividend = self.cs.add(
|
|
ciphertext1=alpha_tensor.dividend, ciphertext2=beta_tensor.dividend
|
|
)
|
|
current_abs_dividend = self.cs.add(
|
|
ciphertext1=alpha_tensor.abs_dividend,
|
|
ciphertext2=beta_tensor.abs_dividend,
|
|
)
|
|
# notice that divisor is alpha tensor's divisor instead of addition
|
|
if alpha_tensor.sign == -1 and beta_tensor.sign == -1:
|
|
current_tensor = Fraction(
|
|
dividend=current_dividend,
|
|
abs_dividend=current_abs_dividend,
|
|
divisor=alpha_tensor.divisor,
|
|
sign=-1,
|
|
)
|
|
else:
|
|
# if one is positive and one is negative, then i cannot know
|
|
# the result is positive or negative. trust mod calculations.
|
|
if alpha_tensor.sign != beta_tensor.sign:
|
|
logger.warn(
|
|
f"{i}-th items of the vectors have different signs, and result's sign "
|
|
"cannot be determined in PHE. Result will be shown for positive for this anyway."
|
|
)
|
|
|
|
current_tensor = Fraction(
|
|
dividend=current_dividend,
|
|
abs_dividend=current_dividend,
|
|
divisor=alpha_tensor.divisor,
|
|
sign=1,
|
|
)
|
|
|
|
current_tensors.append(current_tensor)
|
|
|
|
return EncryptedTensor(fractions=current_tensors, cs=self.cs)
|
|
|
|
|
|
def cast_ciphertext(cs: Homomorphic, value: int) -> Ciphertext:
|
|
"""Cast an integer value to a Ciphertext object."""
|
|
|
|
class_name = cs.__class__.__name__
|
|
algorithm_name = getattr(Algorithm, class_name, None)
|
|
|
|
assert algorithm_name is not None, f"Algorithm name not found for {class_name}"
|
|
|
|
return Ciphertext(
|
|
algorithm_name=algorithm_name,
|
|
keys=cs.keys,
|
|
value=value,
|
|
form=cs.keys.get("form"),
|
|
curve=cs.keys.get("curve"),
|
|
)
|
|
|
|
|
|
def chunkify(lst: list, n: int):
|
|
"""Split list into n approximately equal chunks."""
|
|
avg = len(lst) // n
|
|
remainder = len(lst) % n
|
|
chunks = []
|
|
start = 0
|
|
for i in range(n):
|
|
end = start + avg + (1 if i < remainder else 0)
|
|
chunks.append(lst[start:end])
|
|
start = end
|
|
return chunks
|
|
|
|
|
|
def sum_fractions_chunk(fractions_chunk: list, cs: Homomorphic):
|
|
"""Compute the sum of a chunk of fractions in parallel."""
|
|
result = cast_ciphertext(cs=cs, value=fractions_chunk[0].abs_dividend)
|
|
for fraction in fractions_chunk[1:]:
|
|
result += cast_ciphertext(cs=cs, value=fraction.abs_dividend)
|
|
return result
|