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.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 test_square_inplace(context, plain, precision):
tensor = ts.bfv_tensor(context, plain)
expected = np.square(plain.tolist())
tensor.square_()
decrypted_result = tensor.decrypt().tolist()
assert (decrypted_result == expected).all() == True, "Decryption of tensor is incorrect"
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_power_inplace(context, plain, power, precision):
context = ts.context(ts.SCHEME_TYPE.BFV, 8192, 1032193)
tensor = ts.bfv_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 (decrypted_result == expected) == True, "Decryption of tensor is incorrect"
def test_power(context, plain, power, precision):
context = ts.context(ts.SCHEME_TYPE.BFV, 8192, 1032193)
tensor = ts.bfv_tensor(context, plain)
expected = np.array([np.power(v, power) for v in plain.raw]).reshape(plain.shape).tolist()
new_tensor = tensor ** power
decrypted_result = new_tensor.decrypt().tolist()
assert (decrypted_result == expected) == True, "Decryption of tensor is incorrect"
assert tensor.decrypt().tolist() == plain.tolist(), "Something went wrong in memory."
def test_reshape_batching(context, data, new_shape):
tensor = ts.bfv_tensor(context, data, batch=True)
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_square(context, plain, precision, reshape_first):
tensor = ts.bfv_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 (decrypted_result == expected).all() == True, "Decryption of tensor is incorrect"
assert tensor.decrypt().tolist() == plain.tolist(), "Something went wrong in memory."
def test_negate(context, plain, precision, reshape_first):
tensor = ts.bfv_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 (decrypted_result == expected).all() == True, "Decryption of tensor is incorrect"
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_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 test_polynomial(context, data, polynom, reshape_first):
context = ts.context(ts.SCHEME_TYPE.BFV, 8192, 1032193)
ct = ts.bfv_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 (decrypted_result == expected
) == True, "Polynomial evaluation is incorrect."
def test_sum(context, data, batch, reshape_first, axis, precision):
context.generate_galois_keys()
tensor = ts.bfv_tensor(context, data, batch=batch)
shape = data.shape
full_shape = data.shape
if reshape_first:
shape, full_shape = reshape(batch, shape)
tensor = tensor.reshape(shape)
data = data.reshape(full_shape)
result = tensor.sum(axis)
orig = data.tolist()
np_orig = np.array(orig).reshape(full_shape)
expected = np.sum(np_orig, axis).tolist()
# Decryption
plain_ts = result.decrypt()
decrypted_result = plain_ts.tolist()
assert decrypted_result == expected, "Sum of tensor is incorrect."
assert tensor.decrypt().tolist() == orig, "Something went wrong in memory."
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)
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."
def test_sum_fail(context, data, batch, precision):
context.generate_galois_keys()
tensor = ts.bfv_tensor(context, data, batch=batch)
with pytest.raises(ValueError) as e:
result = tensor.sum(1000)
