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)