def test_lwe_noiseless_trivial_broadcast(thread, src_len):
params = NuFHEParameters()
lwe_size = params.in_out_params.size
res_shape = (10, 20)
src_shape = res_shape[len(res_shape)-src_len:]
res_a = numpy.empty(res_shape + (lwe_size,), Torus32)
res_b = numpy.empty(res_shape, Torus32)
res_cv = numpy.empty(res_shape, ErrorFloat)
mus = get_test_array(src_shape, Torus32)
shape_info = LweSampleArrayShapeInfo(res_a, res_b, res_cv)
test = LweNoiselessTrivial(shape_info, src_shape).compile(thread)
ref = LweNoiselessTrivialReference(shape_info, src_shape)
res_a_dev = thread.empty_like(res_a)
res_b_dev = thread.empty_like(res_b)
res_cv_dev = thread.empty_like(res_cv)
mus_dev = thread.to_device(mus)
test(res_a_dev, res_b_dev, res_cv_dev, mus_dev)
ref(res_a, res_b, res_cv, mus)
assert (res_a_dev.get() == res_a).all()
assert (res_b_dev.get() == res_b).all()
assert errors_allclose(res_cv_dev.get(), res_cv)
def test_lwe_noiseless_trivial_constant(thread):
params = NuFHEParameters()
lwe_size = params.in_out_params.size
shape = (10, 20)
res_a = numpy.empty(shape + (lwe_size,), Torus32)
res_b = numpy.empty(shape, Torus32)
res_cv = numpy.empty(shape, ErrorFloat)
mu = Torus32(-5)
shape_info = LweSampleArrayShapeInfo(res_a, res_b, res_cv)
test = LweNoiselessTrivialConstant(shape_info).compile(thread)
ref = LweNoiselessTrivialConstantReference(shape_info)
res_a_dev = thread.empty_like(res_a)
res_b_dev = thread.empty_like(res_b)
res_cv_dev = thread.empty_like(res_cv)
test(res_a_dev, res_b_dev, res_cv_dev, mu)
ref(res_a, res_b, res_cv, mu)
assert (res_a_dev.get() == res_a).all()
assert (res_b_dev.get() == res_b).all()
assert errors_allclose(res_cv_dev.get(), res_cv)
def test_lwe_keyswitch(thread):
batch_shape = (4, 5)
params = NuFHEParameters()
input_size = params.tgsw_params.tlwe_params.extracted_lweparams.size
output_size = params.in_out_params.size
decomp_length = params.ks_decomp_length
log2_base = params.ks_log2_base
base = 2**log2_base
result_a = numpy.empty(batch_shape + (output_size, ), Torus32)
result_b = numpy.empty(batch_shape, Torus32)
result_cv = numpy.empty(batch_shape, ErrorFloat)
ks_a = get_test_array((input_size, decomp_length, base, output_size),
Torus32, (-1000, 1000))
ks_b = get_test_array((input_size, decomp_length, base), Torus32,
(-1000, 1000))
ks_cv = get_test_array((input_size, decomp_length, base), ErrorFloat,
(-1, 1))
# The base=0 slice of the keyswitch key is a "padding" - it's filled with zeroes.
# The keyswitch function may rely on that.
ks_a[:, :, 0, :] = 0
ks_b[:, :, 0] = 0
ks_cv[:, :, 0] = 0
source_a = get_test_array(batch_shape + (input_size, ), Torus32)
source_b = get_test_array(batch_shape, Torus32, (-1000, 1000))
result_a_dev = thread.empty_like(result_a)
result_b_dev = thread.empty_like(result_b)
result_cv_dev = thread.empty_like(result_cv)
ks_a_dev = thread.to_device(ks_a)
ks_b_dev = thread.to_device(ks_b)
ks_cv_dev = thread.to_device(ks_cv)
source_a_dev = thread.to_device(source_a)
source_b_dev = thread.to_device(source_b)
shape_info = LweSampleArrayShapeInfo(result_a_dev, result_b_dev,
result_cv_dev)
test = LweKeyswitch(shape_info, input_size, output_size, decomp_length,
log2_base).compile(thread)
ref = LweKeyswitchReference(shape_info, input_size, output_size,
decomp_length, log2_base)
test(result_a_dev, result_b_dev, result_cv_dev, ks_a_dev, ks_b_dev,
ks_cv_dev, source_a_dev, source_b_dev)
result_a_test = result_a_dev.get()
result_b_test = result_b_dev.get()
result_cv_test = result_cv_dev.get()
ref(result_a, result_b, result_cv, ks_a, ks_b, ks_cv, source_a, source_b)
assert (result_a == result_a_test).all()
assert (result_b == result_b_test).all()
assert errors_allclose(result_cv, result_cv_test)
def test_lwe_linear_broadcast(thread):
params = NuFHEParameters()
lwe_size = params.in_out_params.size
res_shape = (10, 20)
src_shape = res_shape[1:]
res_a = get_test_array(res_shape + (lwe_size, ), Torus32)
res_b = get_test_array(res_shape, Torus32)
res_cv = get_test_array(res_shape, ErrorFloat, (-1, 1))
src_a = get_test_array(src_shape + (lwe_size, ), Torus32)
src_b = get_test_array(src_shape, Torus32)
src_cv = get_test_array(src_shape, ErrorFloat, (-1, 1))
coeff = 1
add_result = True
res_shape_info = LweSampleArrayShapeInfo(res_a, res_b, res_cv)
src_shape_info = LweSampleArrayShapeInfo(src_a, src_b, src_cv)
test = LweLinear(res_shape_info, src_shape_info,
add_result=add_result).compile(thread)
ref = LweLinearReference(res_shape_info,
src_shape_info,
add_result=add_result)
res_a_dev = thread.to_device(res_a)
res_b_dev = thread.to_device(res_b)
res_cv_dev = thread.to_device(res_cv)
src_a_dev = thread.to_device(src_a)
src_b_dev = thread.to_device(src_b)
src_cv_dev = thread.to_device(src_cv)
thread.synchronize()
test(res_a_dev, res_b_dev, res_cv_dev, src_a_dev, src_b_dev, src_cv_dev,
coeff)
ref(res_a, res_b, res_cv, src_a, src_b, src_cv, coeff)
assert (res_a_dev.get() == res_a).all()
assert (res_b_dev.get() == res_b).all()
assert errors_allclose(res_cv_dev.get(), res_cv)
def test_lwe_linear(thread, positive_coeff, add_result):
params = NuFHEParameters()
lwe_size = params.in_out_params.size
shape = (10, 20)
res_a = get_test_array(shape + (lwe_size, ), Torus32)
res_b = get_test_array(shape, Torus32)
res_cv = get_test_array(shape, Float, (-1, 1))
src_a = get_test_array(shape + (lwe_size, ), Torus32)
src_b = get_test_array(shape, Torus32)
src_cv = get_test_array(shape, Float, (-1, 1))
coeff = 1 if positive_coeff else -1
shape_info = LweSampleArrayShapeInfo(src_a, src_b, src_cv)
test = LweLinear(shape_info, shape_info,
add_result=add_result).compile(thread)
ref = LweLinearReference(shape_info, shape_info, add_result=add_result)
res_a_dev = thread.to_device(res_a)
res_b_dev = thread.to_device(res_b)
res_cv_dev = thread.to_device(res_cv)
src_a_dev = thread.to_device(src_a)
src_b_dev = thread.to_device(src_b)
src_cv_dev = thread.to_device(src_cv)
thread.synchronize()
test(res_a_dev, res_b_dev, res_cv_dev, src_a_dev, src_b_dev, src_cv_dev,
coeff)
ref(res_a, res_b, res_cv, src_a, src_b, src_cv, coeff)
assert (res_a_dev.get() == res_a).all()
assert (res_b_dev.get() == res_b).all()
assert numpy.allclose(res_cv_dev.get(), res_cv)