def test_matmul_api(context, plain, arithmetic):
r_t = np.random.randint(0, 100, size=(2, 2), dtype=np.int64)
l_t = np.random.randint(0, 100, size=(2, 2), dtype=np.int64)
r_pt = ts.plain_tensor(r_t.flatten().tolist(), (2, 2), dtype="int")
l_pt = ts.plain_tensor(l_t.flatten().tolist(), (2, 2), dtype="int")
right = ts.bfv_tensor(context, r_pt)
if plain:
left = l_pt
else:
left = ts.bfv_tensor(context, l_pt)
expected_result = r_t.dot(l_t)
## non-inplace
if arithmetic:
result = right @ left
else:
result = right.mm(left)
np_result = np.array(result.decrypt().tolist())
assert np_result.shape == expected_result.shape
assert np.allclose(np_result, expected_result, rtol=0, atol=0)
# right didn't change
right_result = np.array(right.decrypt().tolist())
assert np.allclose(right_result, r_t, rtol=0, atol=0)
# left didn't change
if plain:
left_result = l_t
else:
left_result = np.array(left.decrypt().tolist())
assert np.allclose(left_result, l_t, rtol=0, atol=0)
# inplace
if arithmetic:
right @= left
else:
right.mm_(left)
np_result = np.array(result.decrypt().tolist())
assert np_result.shape == expected_result.shape
assert np.allclose(np_result, expected_result, rtol=0, atol=0)
# right didn't change
right_result = np.array(right.decrypt().tolist())
assert right_result.shape == expected_result.shape
assert np.allclose(right_result, expected_result, rtol=0, atol=0)
# left didn't change
if plain:
left_result = l_t
else:
left_result = np.array(left.decrypt().tolist())
assert np.allclose(left_result, l_t, rtol=0, atol=0)
def test_dot(context, shapes, plain):
r_shape = shapes[0]
l_shape = shapes[1]
r_t = np.random.randn(*r_shape)
l_t = np.random.randn(*l_shape)
r_pt = ts.plain_tensor(r_t.flatten().tolist(), r_shape)
l_pt = ts.plain_tensor(l_t.flatten().tolist(), l_shape)
right = ts.ckks_tensor(context, r_pt)
if plain:
left = l_pt
else:
left = ts.ckks_tensor(context, l_pt)
expected_result = r_t.dot(l_t)
## non-inplace
result = right.dot(left)
np_result = np.array(result.decrypt().tolist())
assert np_result.shape == expected_result.shape
assert np.allclose(np_result, expected_result, rtol=0, atol=0.01)
# right didn't change
right_result = np.array(right.decrypt().tolist())
assert np.allclose(right_result, r_t, rtol=0, atol=0.01)
# left didn't change
if plain:
left_result = l_t
else:
left_result = np.array(left.decrypt().tolist())
assert np.allclose(left_result, l_t, rtol=0, atol=0.01)
# inplace
right.dot_(left)
np_result = np.array(result.decrypt().tolist())
assert np_result.shape == expected_result.shape
assert np.allclose(np_result, expected_result, rtol=0, atol=0.01)
# right didn't change
right_result = np.array(right.decrypt().tolist())
assert right_result.shape == expected_result.shape
assert np.allclose(right_result, expected_result, rtol=0, atol=0.01)
# left didn't change
if plain:
left_result = l_t
else:
left_result = np.array(left.decrypt().tolist())
assert np.allclose(left_result, l_t, rtol=0, atol=0.01)
def test_dot(context, shapes, plain):
r_shape = shapes[0]
l_shape = shapes[1]
r_t = np.random.randint(0, 100, *[r_shape], dtype=np.int64)
l_t = np.random.randint(0, 100, *[l_shape], dtype=np.int64)
r_pt = ts.plain_tensor(r_t.flatten().tolist(), r_shape, dtype="int")
l_pt = ts.plain_tensor(l_t.flatten().tolist(), l_shape, dtype="int")
right = ts.bfv_tensor(context, r_pt)
if plain:
left = l_pt
else:
left = ts.bfv_tensor(context, l_pt)
expected_result = r_t.dot(l_t)
## non-inplace
result = right.dot(left)
np_result = np.array(result.decrypt().tolist())
assert np_result.shape == expected_result.shape
assert np.allclose(np_result, expected_result, rtol=0, atol=0)
# right didn't change
right_result = np.array(right.decrypt().tolist())
assert np.allclose(right_result, r_t, rtol=0, atol=0)
# left didn't change
if plain:
left_result = l_t
else:
left_result = np.array(left.decrypt().tolist())
assert np.allclose(left_result, l_t, rtol=0, atol=0)
# inplace
right.dot_(left)
np_result = np.array(result.decrypt().tolist())
assert np_result.shape == expected_result.shape
assert np.allclose(np_result, expected_result, rtol=0, atol=0)
# right didn't change
right_result = np.array(right.decrypt().tolist())
assert right_result.shape == expected_result.shape
assert np.allclose(right_result, expected_result, rtol=0, atol=0)
# left didn't change
if plain:
left_result = l_t
else:
left_result = np.array(left.decrypt().tolist())
assert np.allclose(left_result, l_t, rtol=0, atol=0)
def __init__(self,
context: "ts.Context" = None,
vector=None,
data: ts._ts_cpp.BFVVector = None):
"""Constructor method for the BFVVector object, which can store a vector of
integers in encrypted form, using the BFV homomorphic encryption scheme.
Args:
context: a Context object, holding the encryption parameters and keys.
vector (of int): a vector holding data to be encrypted.
data: A ts._ts_cpp.BFVVector to wrap. We won't construct a new object if it's passed.
Returns:
BFVVector object.
"""
# wrapping
if data is not None:
self.data = data
# constructing a new object
else:
if not isinstance(context, ts.Context):
raise TypeError("context must be a tenseal.Context")
if not isinstance(vector, ts.PlainTensor):
vector = ts.plain_tensor(vector, dtype="int")
if len(vector.shape) != 1:
raise ValueError("can only encrypt a vector")
vector = vector.raw
self.data = ts._ts_cpp.BFVVector(context.data, vector)
def __init__(
self,
context: "ts.Context" = None,
tensor=None,
batch: bool = False,
data: ts._ts_cpp.BFVTensor = None,
):
"""Constructor method for the BFVTensor object, which can store an n-dimensional
int tensor in encrypted form, using the BFV homomorphic encryption scheme.
Args:
context: a Context object, holding the encryption parameters and keys.
tensor: tensor-like object.
batch: should we use ciphertext-level batching?
data: A ts._ts_cpp.BFVTensor to wrap. We won't construct a new object if it's passed.
Returns:
BFVTensor object.
"""
# wrapping
if data is not None:
self.data = data
# constructing a new object
else:
if not isinstance(context, ts.Context):
raise TypeError("context must be a tenseal.Context")
if not isinstance(tensor, ts.PlainTensor):
tensor = ts.plain_tensor(tensor, dtype="int")
tensor.dtype = "int"
self.data = ts._ts_cpp.BFVTensor(context.data, tensor.data, batch)
def test_sanity(data, shape):
tensor = ts.plain_tensor(data, shape)
strides = []
stride = 1
for s in reversed(shape):
strides.append(stride)
stride *= s
strides = list(reversed(strides))
assert tensor.raw == data
assert tensor.shape == shape
assert tensor.size() == shape[0]
assert len(tensor) == shape[0]
assert tensor.empty() == False
assert tensor.strides() == strides
buf = tensor.serialize()
new_tensor = ts.plain_tensor_from(buf)
assert new_tensor.raw == data
assert new_tensor.shape == shape
assert new_tensor.size() == shape[0]
assert len(new_tensor) == shape[0]
assert new_tensor.empty() == False
assert new_tensor.strides() == strides
def test_broadcast(data, shape, new_shape):
tensor = ts.plain_tensor(data, shape)
newt = tensor.broadcast(new_shape)
assert tensor.shape == shape
assert newt.shape == new_shape
tensor.broadcast_(new_shape)
assert tensor.shape == new_shape
def test_broadcast(context, data, shape, new_shape):
tensor = ts.bfv_tensor(context, ts.plain_tensor(data, shape, dtype="int"))
newt = tensor.broadcast(new_shape)
assert tensor.shape == shape
assert newt.shape == new_shape
tensor.broadcast_(new_shape)
assert tensor.shape == new_shape
def test_ckks_tensor_sanity(plain_vec, precision, duplicate):
context = ckks_context()
plain_tensor = ts.plain_tensor(plain_vec)
orig = ts.ckks_tensor(context, plain_tensor)
ckks_tensor = duplicate(orig)
decrypted = ckks_tensor.decrypt().tolist()
assert _almost_equal(decrypted, plain_vec,
precision), "Decryption of tensor is incorrect"
def test_reshape(data, shape, reshape):
tensor = ts.plain_tensor(data, shape)
newt = tensor.reshape(reshape)
assert tensor.shape == shape
assert newt.shape == reshape
tensor.reshape_(reshape)
assert tensor.shape == reshape
def _enc_matmul_plain(self, other):
if not isinstance(other, ts.PlainTensor):
try:
other = ts.plain_tensor(other, dtype="float")
except TypeError:
raise TypeError(f"can't operate with object of type {type(other)}")
if len(other.shape) != 1:
raise ValueError("can only operate with a vector")
other = other.raw
return other
def test_dot_product_plain(context, vec1, vec2):
context.generate_galois_keys()
first_vec = ts.bfv_vector(context, vec1)
second_vec = ts.plain_tensor(vec2, dtype="int")
result = first_vec.dot(second_vec)
expected = [sum([v1 * v2 for v1, v2 in zip(vec1, vec2)])]
# Decryption
assert result.decrypt() == expected, "Dot product of vectors is incorrect."
assert first_vec.decrypt() == vec1, "Something went wrong in memory."
def test_dot_product_plain_inplace(context, vec1, vec2):
context.generate_galois_keys()
first_vec = ts.bfv_vector(context, vec1)
second_vec = ts.plain_tensor(vec2, dtype="int")
first_vec.dot_(second_vec)
expected = [sum([v1 * v2 for v1, v2 in zip(vec1, vec2)])]
# Decryption
assert first_vec.decrypt(
) == expected, "Dot product of vectors is incorrect."
def _mm(cls, other):
if not isinstance(other, ts.PlainTensor):
try:
other = ts.plain_tensor(other, dtype="float")
except TypeError:
raise TypeError(f"can't operate with object of type {type(other)}")
if len(other.shape) != 2:
raise ValueError("can only operate with a matrix")
other = other.tolist()
return other
def test_ckks_tensor_lazy_load(precision):
vec1 = [1, 2, 3, 4]
vec2 = [1, 2, 3, 4]
context = ckks_context()
first_vec = ts.ckks_tensor(context, ts.plain_tensor(vec1))
second_vec = ts.ckks_tensor(context, ts.plain_tensor(vec2))
buff = first_vec.serialize()
newvec = ts.lazy_ckks_tensor_from(buff)
newvec.link_context(context)
result = newvec + second_vec
# Decryption
decrypted_result = result.decrypt().tolist()
assert _almost_equal(decrypted_result, [2, 4, 6, 8],
precision), "Decryption of tensor is incorrect"
assert _almost_equal(newvec.decrypt().tolist(), [1, 2, 3, 4],
precision), "invalid new tensor"
def test_dot_product_plain_inplace(context, vec1, vec2, precision):
context.generate_galois_keys()
first_vec = ts.ckks_vector(context, vec1)
second_vec = ts.plain_tensor(vec2)
first_vec.dot_(second_vec)
expected = [sum([v1 * v2 for v1, v2 in zip(vec1, vec2)])]
# Decryption
decrypted_result = first_vec.decrypt()
assert _almost_equal(decrypted_result, expected,
precision), "Dot product of vectors is incorrect."
def _dot(cls, other):
if isinstance(other, (cls)):
return other.data
if not isinstance(other, ts.PlainTensor):
try:
other = ts.plain_tensor(other, dtype="float")
except TypeError:
raise TypeError(f"can't operate with object of type {type(other)}")
if len(other.shape) != 1:
raise ValueError("can only operate with a vector")
return other.data
def _get_operand(cls, other, dtype: str = "float") -> Union[int, float, "ts._ts_cpp.Tensor"]:
"""Extract the appropriate operand the tensor can operate with"""
if isinstance(other, (int, float)):
return other
elif isinstance(other, (cls, ts.PlainTensor)):
return other.data
else:
try:
other = ts.plain_tensor(other, dtype=dtype)
other = other.data
except TypeError:
raise TypeError(f"can't operate with object of type {type(other)}")
return other
def test_ckks_tensor_encryption_decryption(batch, shape):
context = ts.context(ts.SCHEME_TYPE.CKKS,
8192,
coeff_mod_bit_sizes=COEFF_MOD_BIT_SIZES)
scale = pow(2, 40)
tensor = np.random.randn(*shape)
plain_tensor = ts.plain_tensor(tensor.flatten().tolist(), shape=shape)
ckks_vec = ts.ckks_tensor(context, plain_tensor, scale, batch)
decrypted_vec = ckks_vec.decrypt().tolist()
assert np.array(decrypted_vec).shape == tensor.shape
assert np.allclose(tensor, decrypted_vec, rtol=0, atol=0.001)
def test_transpose(data, shape):
tensor = ts.plain_tensor(data, shape)
expected = np.transpose(np.array(data).reshape(shape))
newt = tensor.transpose()
assert tensor.shape == shape
assert newt.shape == list(expected.shape)
assert np.array(newt.tolist()).any() == expected.any()
tensor.transpose_()
assert tensor.shape == list(expected.shape)
assert np.array(tensor.tolist()).any() == expected.any()
def test_transpose(context, data, shape):
tensor = ts.bfv_tensor(context, ts.plain_tensor(data, shape, dtype="int"))
expected = np.transpose(np.array(data).reshape(shape))
newt = tensor.transpose()
assert tensor.shape == shape
assert newt.shape == list(expected.shape)
result = np.array(newt.decrypt().tolist())
assert np.allclose(result, expected, rtol=0, atol=0)
tensor.transpose_()
assert tensor.shape == list(expected.shape)
result = np.array(tensor.decrypt().tolist())
assert np.allclose(result, expected, rtol=0, atol=0)
def test_ckks_tensor_encryption_decryption_matrix(batch):
context = ts.context(ts.SCHEME_TYPE.CKKS,
8192,
coeff_mod_bit_sizes=COEFF_MOD_BIT_SIZES)
scale = pow(2, 40)
matrix = np.random.randn(4, 5)
plain_tensor = ts.plain_tensor(matrix.tolist())
ckks_vec = ts.ckks_tensor(context, plain_tensor, scale, batch)
decrypted_vec = ckks_vec.decrypt().tolist()
assert len(decrypted_vec) == len(matrix)
for idx in range(len(decrypted_vec)):
row = decrypted_vec[idx]
assert isinstance(row, list)
assert len(row) == len(matrix[0])
assert _almost_equal(row, matrix[idx],
1), "Decryption of vector is incorrect"
def test_add_sub_mul_scalar(context, shape, op):
r_t = np.random.randn(*shape)
r_pt = ts.plain_tensor(r_t.flatten().tolist(), shape)
right = ts.ckks_tensor(context, r_pt)
left = np.random.randn(1)[0]
if op == "add":
expected_result = r_t + left
elif op == "sub":
expected_result = r_t - left
elif op == "mul":
expected_result = r_t * left
## non-inplace
if op == "add":
result = right + left
elif op == "sub":
result = right - left
elif op == "mul":
result = right * left
np_result = np.array(result.decrypt().tolist())
assert np_result.shape == expected_result.shape
assert np.allclose(np_result, expected_result, rtol=0, atol=0.01)
# right didn't change
right_result = np.array(right.decrypt().tolist())
assert np.allclose(right_result, r_t, rtol=0, atol=0.01)
# inplace
if op == "add":
right += left
elif op == "sub":
right -= left
elif op == "mul":
right *= left
np_result = np.array(result.decrypt().tolist())
assert np_result.shape == expected_result.shape
assert np.allclose(np_result, expected_result, rtol=0, atol=0.01)
# right didn't change
right_result = np.array(right.decrypt().tolist())
assert right_result.shape == expected_result.shape
assert np.allclose(right_result, expected_result, rtol=0, atol=0.01)
def test_add_sub_mul_scalar(context, shape, op):
r_t = np.random.randint(0, 100, *[shape], dtype=np.int64)
r_pt = ts.plain_tensor(r_t.flatten().tolist(), shape, dtype="int")
right = ts.bfv_tensor(context, r_pt)
left = np.random.randint(0, 100, size=1, dtype=np.int64)[0]
if op == "add":
expected_result = r_t + left
elif op == "sub":
expected_result = r_t - left
elif op == "mul":
expected_result = r_t * left
## non-inplace
if op == "add":
result = right + left
elif op == "sub":
result = right - left
elif op == "mul":
result = right * left
np_result = np.array(result.decrypt().tolist())
assert np_result.shape == expected_result.shape
assert np.allclose(np_result, expected_result, rtol=0, atol=0)
# right didn't change
right_result = np.array(right.decrypt().tolist())
assert np.allclose(right_result, r_t, rtol=0, atol=0)
# inplace
if op == "add":
right += left
elif op == "sub":
right -= left
elif op == "mul":
right *= left
np_result = np.array(result.decrypt().tolist())
assert np_result.shape == expected_result.shape
assert np.allclose(np_result, expected_result, rtol=0, atol=0)
# right didn't change
right_result = np.array(right.decrypt().tolist())
assert right_result.shape == expected_result.shape
assert np.allclose(right_result, expected_result, rtol=0, atol=0)
def __init__(
self,
context: "ts.Context" = None,
tensor=None,
scale: float = None,
batch: bool = False,
data: ts._ts_cpp.CKKSTensor = None,
):
"""Constructor method for the CKKSTensor object, which can store an n-dimensional
float tensor in encrypted form, using the CKKS homomorphic encryption scheme.
Args:
context: a Context object, holding the encryption parameters and keys.
tensor: tensor-like object.
scale: the scale to be used to encode tensor values. CKKSTensor will use the global_scale provided by the context if it's set to None.
batch: should we use ciphertext-level batching?
data: A ts._ts_cpp.CKKSTensor to wrap. We won't construct a new object if it's passed.
Returns:
CKKSTensor object.
"""
# wrapping
if data is not None:
self.data = data
# constructing a new object
else:
if not isinstance(context, ts.Context):
raise TypeError("context must be a tenseal.Context")
if not isinstance(tensor, ts.PlainTensor):
tensor = ts.plain_tensor(tensor, dtype="float")
tensor.dtype = "float"
if scale is None:
self.data = ts._ts_cpp.CKKSTensor(context.data, tensor.data,
batch)
else:
self.data = ts._ts_cpp.CKKSTensor(context.data, tensor.data,
scale, batch)
def __init__(
self,
context: "ts.Context" = None,
vector=None,
scale: float = None,
data: ts._ts_cpp.CKKSVector = None,
):
"""Constructor method for the CKKSVector object, which can store a vector of
float numbers in encrypted form, using the CKKS homomorphic encryption scheme.
Args:
context: a Context object, holding the encryption parameters and keys.
vector (of float): a vector holding data to be encrypted.
scale: the scale to be used to encode vector values. CKKSTensor will use the global_scale provided by the context if it's set to None.
data: A ts._ts_cpp.CKKSVector to wrap. We won't construct a new object if it's passed.
Returns:
CKKSVector object.
"""
# wrapping
if data is not None:
self.data = data
# constructing a new object
else:
if not isinstance(context, ts.Context):
raise TypeError("context must be a tenseal.Context")
if not isinstance(vector, ts.PlainTensor):
vector = ts.plain_tensor(vector, dtype="float")
if len(vector.shape) != 1:
raise ValueError("can only encrypt a vector")
vector = vector.raw
if scale is None:
self.data = ts._ts_cpp.CKKSVector(context.data, vector)
else:
self.data = ts._ts_cpp.CKKSVector(context.data, vector, scale)
def test_broadcast_add_sub_mul_tensor_ct_pt(context, shape, plain, op):
l_t = np.random.randint(0, 100, *[shape[1]], dtype=np.int64)
r_t = np.random.randint(0, 100, *[shape[0]], dtype=np.int64)
l_pt = ts.plain_tensor(l_t.flatten().tolist(), shape[1], dtype="int")
r_pt = ts.plain_tensor(r_t.flatten().tolist(), shape[0], dtype="int")
right = ts.bfv_tensor(context, r_pt)
if plain:
left = l_pt
else:
left = ts.bfv_tensor(context, l_pt)
if op == "add":
expected_result = r_t + l_t
elif op == "sub":
expected_result = r_t - l_t
elif op == "mul":
expected_result = r_t * l_t
## non-inplace
if op == "add":
result = right + left
elif op == "sub":
result = right - left
elif op == "mul":
result = right * left
np_result = np.array(result.decrypt().tolist())
assert np_result.shape == expected_result.shape
assert np.allclose(np_result, expected_result, rtol=0, atol=0)
# right didn't change
right_result = np.array(right.decrypt().tolist())
assert np.allclose(right_result, r_t, rtol=0, atol=0)
# left didn't change
if plain:
left_result = l_t
else:
left_result = np.array(left.decrypt().tolist())
assert np.allclose(left_result, l_t, rtol=0, atol=0)
# inplace
if op == "add":
right += left
elif op == "sub":
right -= left
elif op == "mul":
right *= left
np_result = np.array(result.decrypt().tolist())
assert np_result.shape == expected_result.shape
assert np.allclose(np_result, expected_result, rtol=0, atol=0)
# right didn't change
right_result = np.array(right.decrypt().tolist())
assert right_result.shape == expected_result.shape
assert np.allclose(right_result, expected_result, rtol=0, atol=0)
# left didn't change
if plain:
left_result = l_t
else:
left_result = np.array(left.decrypt().tolist())
assert np.allclose(left_result, l_t, rtol=0, atol=0)
batch_size = new_shape[0]
new_shape = new_shape[1:]
new_shape = new_shape[-1:] + new_shape[:-1]
full_shape = new_shape.copy()
if batched:
full_shape.insert(0, batch_size)
return new_shape, full_shape
@pytest.mark.parametrize(
"data, new_shape",
[
# (ts.plain_tensor([0]), 0),
(ts.plain_tensor([i for i in range(8)], dtype="int"), [4, 2]),
(ts.plain_tensor([i for i in range(210)],
shape=[2, 3, 5, 7],
dtype="int"), [7, 15, 2]),
],
)
def test_reshape_no_batching(context, data, new_shape):
tensor = ts.bfv_tensor(context, data, batch=False)
old_shape = tensor.shape
new_t = tensor.reshape(new_shape)
assert new_t.shape == new_shape
assert tensor.shape == old_shape
tensor.reshape_(new_shape)
)
def test_reshape(data, shape, reshape):
tensor = ts.plain_tensor(data, shape)
newt = tensor.reshape(reshape)
assert tensor.shape == shape
assert newt.shape == reshape
tensor.reshape_(reshape)
assert tensor.shape == reshape
@pytest.mark.parametrize(
"data, axis",
[
(ts.plain_tensor([0]), 0),
(ts.plain_tensor([i for i in range(8)]), 0),
(ts.plain_tensor([i for i in range(6)], shape=[2, 3]), 0),
(ts.plain_tensor([i for i in range(6)], shape=[2, 3]), 1),
(ts.plain_tensor([i for i in range(8)], shape=[2, 2, 2]), 1),
(ts.plain_tensor([i for i in range(30)], shape=[2, 3, 5]), 0),
(ts.plain_tensor([i for i in range(30)], shape=[2, 3, 5]), 1),
(ts.plain_tensor([i for i in range(30)], shape=[2, 3, 5]), 2),
(ts.plain_tensor([i for i in range(210)], shape=[2, 3, 5, 7]), 0),
(ts.plain_tensor([i for i in range(210)], shape=[2, 3, 5, 7]), 1),
(ts.plain_tensor([i for i in range(210)], shape=[2, 3, 5, 7]), 2),
(ts.plain_tensor([i for i in range(210)], shape=[2, 3, 5, 7]), 3),
],
)
def test_batch(data, axis):
batch = data.batch(axis)
def test_size(context):
for size in range(1, 10):
vec = ts.bfv_tensor(context, ts.plain_tensor([1] * size, dtype="int"))
assert vec.shape == [size], "Size of encrypted tensor is incorrect."
