def test_matmul_api(context, plain, arithmetic):
r_t = np.random.randn(2, 2)
l_t = np.random.randn(2, 2)
r_pt = ts.plain_tensor(r_t.flatten().tolist(), (2, 2))
l_pt = ts.plain_tensor(l_t.flatten().tolist(), (2, 2))
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
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.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
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.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.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_sum_batch(context, data, reshape_first, precision):
context.generate_galois_keys()
tensor = ts.ckks_tensor(context, data, batch=True)
shape = data.shape
full_shape = data.shape
if reshape_first:
shape, full_shape = reshape(True, shape)
tensor = tensor.reshape(shape)
data = data.reshape(full_shape)
orig = data.tolist()
np_orig = np.array(orig).reshape(full_shape)
expected = np.sum(np_orig, 0).tolist()
result = tensor.sum_batch()
# Decryption
plain_ts = result.decrypt()
decrypted_result = plain_ts.tolist()
assert _almost_equal(decrypted_result, expected,
precision), "Sum of tensor is incorrect."
assert _almost_equal(tensor.decrypt().tolist(), orig,
precision), "Something went wrong in memory."
def test_square_inplace(context, plain, precision):
tensor = ts.ckks_tensor(context, plain)
expected = np.square(plain.tolist())
tensor.square_()
decrypted_result = tensor.decrypt().tolist()
assert _almost_equal(decrypted_result, expected,
precision), "Decryption of tensor is incorrect"
def test_polynomial_rescale_off(context, data, polynom):
context = ts.context(ts.SCHEME_TYPE.CKKS, 8192, 0, [60, 40, 40, 60])
context.global_scale = 2**40
context.auto_rescale = False
ct = ts.ckks_tensor(context, data)
with pytest.raises(ValueError) as e:
result = ct.polyval(polynom)
def test_broadcast(context, data, shape, new_shape):
tensor = ts.ckks_tensor(context, 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_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 __encrypt(self, context, state_dict):
if not isinstance(context, ts.Context):
raise TypeError('Invalid input types context: {}'.format(
type(context)))
aux_state_dict = state_dict.copy()
for name, tensor in aux_state_dict.items():
aux_state_dict[name] = ts.ckks_tensor(context, tensor)
return aux_state_dict
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_reshape_no_batching(context, data, new_shape):
tensor = ts.ckks_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)
assert tensor.shape == new_shape
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_subscript(context, data, slices, new_shape):
tensor = ts.ckks_tensor(context, data)
plain_data = np.array(data)
new_tensor = tensor[slices]
assert new_tensor.shape == new_shape
if isinstance(slices, int):
assert _almost_equal(new_tensor.decrypt().tolist(),
[plain_data.__getitem__(slices).tolist()], 1)
else:
assert _almost_equal(new_tensor.decrypt().tolist(),
plain_data.__getitem__(slices).tolist(), 1)
def test_negate(context, plain, precision, reshape_first):
tensor = ts.ckks_tensor(context, plain)
if reshape_first:
shape, _ = reshape(False, plain.shape)
tensor = tensor.reshape(shape)
plain = plain.reshape(shape)
expected = np.negative(plain.tolist())
result = -tensor
decrypted_result = result.decrypt().tolist()
assert _almost_equal(decrypted_result, expected,
precision), "Decryption of tensor is incorrect"
def test_power_inplace(context, plain, power, precision):
context = ts.context(ts.SCHEME_TYPE.CKKS,
16384,
coeff_mod_bit_sizes=[60, 40, 40, 40, 40, 60])
context.global_scale = pow(2, 40)
tensor = ts.ckks_tensor(context, plain)
expected = np.array([np.power(v, power)
for v in plain.raw]).reshape(plain.shape).tolist()
tensor **= power
decrypted_result = tensor.decrypt().tolist()
assert _almost_equal(decrypted_result, expected,
precision), "Decryption of tensor is incorrect"
def test_transpose(context, data, shape):
tensor = ts.ckks_tensor(context, 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)
result = np.array(newt.decrypt().tolist())
assert np.allclose(result, expected, rtol=0, atol=0.01)
tensor.transpose_()
assert tensor.shape == list(expected.shape)
result = np.array(tensor.decrypt().tolist())
assert np.allclose(result, expected, rtol=0, atol=0.01)
def test_square(context, plain, precision, reshape_first):
tensor = ts.ckks_tensor(context, plain)
result = tensor.square()
if reshape_first:
shape, _ = reshape(False, plain.shape)
result = result.reshape(shape)
plain = plain.reshape(shape)
expected = np.square(plain.tolist())
decrypted_result = result.decrypt().tolist()
assert _almost_equal(decrypted_result, expected,
precision), "Decryption of tensor is incorrect"
assert _almost_equal(tensor.decrypt().tolist(), plain.tolist(),
precision), "Something went wrong in memory."
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_polynomial(context, data, polynom, reshape_first):
ct = ts.ckks_tensor(context, data)
shape = data.shape
if reshape_first:
shape, _ = reshape(False, shape)
data = data.reshape(shape)
expected = (np.array([np.polyval(polynom[::-1], x)
for x in data.raw]).reshape(data.shape).tolist())
result = ct.polyval(polynom)
result = result.reshape(shape)
if len(polynom) >= 13:
precision = -1
else:
precision = 1
decrypted_result = result.decrypt().tolist()
assert _almost_equal(decrypted_result, expected,
precision), "Polynomial evaluation is incorrect."
def test_sum_fail(context, data, batch, precision):
context.generate_galois_keys()
tensor = ts.ckks_tensor(context, data, batch=batch)
with pytest.raises(ValueError) as e:
result = tensor.sum(1000)
def test_size(context):
for size in range(1, 10):
vec = ts.ckks_tensor(context, ts.plain_tensor([1] * size))
assert vec.shape == [size], "Size of encrypted tensor is incorrect."
def test_broadcast_add_sub_mul_tensor_ct_pt(context, shape, plain, op):
l_t = np.random.randn(*shape[1])
r_t = np.random.randn(*shape[0])
l_pt = ts.plain_tensor(l_t.flatten().tolist(), shape[1])
r_pt = ts.plain_tensor(r_t.flatten().tolist(), shape[0])
right = ts.ckks_tensor(context, r_pt)
if plain:
left = l_pt
else:
left = ts.ckks_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.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
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)
# 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)
