def test_primitive_logging_model(get_clients) -> None:
model = LinearNet(torch)
clients = get_clients(2)
session = Session(parties=clients)
SessionManager.setup_mpc(session)
mpc_model = model.share(session=session)
x_secret = torch.randn(2, 3)
x_mpc = MPCTensor(secret=x_secret, session=session)
model.eval()
expected_primitive_log = {
"beaver_matmul": [(
{
"a_shape": (2, 3),
"b_shape": (3, 10)
},
{
"a_shape": (2, 3),
"b_shape": (3, 10),
"nr_parties": 2
},
)],
"fss_comp": [({}, {
"n_values": 20
})],
"beaver_mul": [(
{
"a_shape": (2, 10),
"b_shape": (2, 10)
},
{
"a_shape": (2, 10),
"b_shape": (2, 10),
"nr_parties": 2
},
)],
}
CryptoPrimitiveProvider.start_logging()
mpc_model(x_mpc)
primitive_log = CryptoPrimitiveProvider.stop_logging()
assert expected_primitive_log == primitive_log
def test_ops_integer(get_clients, nr_clients, op_str) -> None:
clients = get_clients(nr_clients)
session = Session(parties=clients)
SessionManager.setup_mpc(session)
op = getattr(operator, op_str)
x_secret = torch.Tensor([0.125, -1.25, -4.25, 4])
y = 4
x = MPCTensor(secret=x_secret, session=session)
expected_result = op(x_secret, y)
result = op(x, y).reconstruct()
assert np.allclose(result, expected_result, atol=10e-3)
def test_private_matmul(get_clients, security):
parties = get_clients(3)
protocol = Falcon(security)
session = Session(protocol=protocol, parties=parties)
SessionManager.setup_mpc(session)
secret1 = torch.tensor([
[[-100.25, 20.3, 30.12], [-50.1, 100.217, 1.2],
[1032.15, -323.56, 15.15]],
[[-0.25, 2.3, 3012], [-5.01, 1.00217, 1.2], [2.15, -3.56, 15.15]],
])
secret2 = torch.tensor([[-1, 0.28, 3], [-9, 10.18, 1], [32, -23, 5]])
tensor1 = MPCTensor(secret=secret1, session=session)
tensor2 = MPCTensor(secret=secret2, session=session)
result = tensor1 @ tensor2
expected_res = secret1 @ secret2
assert np.allclose(result.reconstruct(), expected_res, atol=1e-3)
def test_local_secret_not_tensor(get_clients) -> None:
alice_client, bob_client = get_clients(2)
session = Session(parties=[alice_client, bob_client])
SessionManager.setup_mpc(session)
x_int = 5
x = MPCTensor(secret=x_int, session=session)
result = x.reconstruct()
assert x_int == result
x_float = 5.987
x = MPCTensor(secret=x_float, session=session)
result = x.reconstruct()
assert np.allclose(torch.tensor(x_float), result)
def test_remote_not_tensor(get_clients) -> None:
alice_client, bob_client = get_clients(2)
session = Session(parties=[alice_client, bob_client])
SessionManager.setup_mpc(session)
x_remote_int = bob_client.python.Int(5)
x = MPCTensor(secret=x_remote_int, shape=(1, ), session=session)
result = x.reconstruct()
assert x_remote_int == result
x_remote_int = bob_client.python.Float(5.4)
x = MPCTensor(secret=x_remote_int, shape=(1, ), session=session)
result = x.reconstruct()
assert np.allclose(x_remote_int.get(), result, atol=1e-5)
def test_ops_mpc_private_rst(get_clients, nr_clients, op_str) -> None:
clients = get_clients(nr_clients)
falcon = Protocol.registered_protocols["Falcon"]()
session = Session(parties=clients, protocol=falcon)
SessionManager.setup_mpc(session)
op = getattr(operator, op_str)
x_secret = torch.Tensor([[0.125, -1.25], [-4.25, 4]])
y_secret = torch.Tensor([[4.5, -2.5], [5, 2.25]])
x = MPCTensor(secret=x_secret, session=session)
y = MPCTensor(secret=y_secret, session=session)
expected_result = op(x_secret, y_secret)
result = op(x, y).reconstruct()
assert np.allclose(result, expected_result, atol=10e-4)
def test_reconstruct(get_clients) -> None:
clients = get_clients(2)
session = Session(parties=clients)
SessionManager.setup_mpc(session)
a_rand = 3
a = ShareTensor(data=a_rand, encoder_precision=0)
a_shares = MPCTensor.generate_shares(a, 2, torch.long)
a_shares_copy = MPCTensor.generate_shares(a_rand, 2, torch.long)
x_secret = torch.Tensor([1, -2, 3.0907, -4.870])
x = MPCTensor(secret=x_secret, session=session)
x = x.reconstruct()
assert torch.allclose(x_secret, x)
def test_ops_public_mul(get_clients, security, base, precision):
parties = get_clients(3)
protocol = Falcon(security)
config = Config(encoder_base=base, encoder_precision=precision)
session = Session(protocol=protocol, parties=parties, config=config)
SessionManager.setup_mpc(session)
secret = torch.Tensor([[0.125, 1001, -1.25, 4.82], [-4.25, 0.217, 3301,
4]])
value = 8
tensor = MPCTensor(secret=secret, session=session)
result = tensor * value
expected_res = secret * value
assert np.allclose(result.reconstruct(), expected_res, atol=1e-3)
def test_mul_private(get_clients, security, base, precision):
parties = get_clients(3)
protocol = Falcon(security)
config = Config(encoder_base=base, encoder_precision=precision)
session = Session(protocol=protocol, parties=parties, config=config)
SessionManager.setup_mpc(session)
secret1 = torch.tensor([[-100.25, 0.29, 30.45], [-90.82, 1000, 0.18],
[1032.45, -323.18, 15.15]])
secret2 = 8
tensor1 = MPCTensor(secret=secret1, session=session)
tensor2 = MPCTensor(secret=secret2, session=session)
result = tensor1 * tensor2
expected_res = secret1 * secret2
assert np.allclose(result.reconstruct(), expected_res, atol=1e-3)
def test_mse_loss(get_clients, reduction) -> None:
clients = get_clients(4)
session = Session(parties=clients)
SessionManager.setup_mpc(session)
y_secret = torch.Tensor([0.23, 0.32, 0.2, 0.3])
y_mpc = MPCTensor(secret=y_secret, session=session)
y_pred = torch.Tensor([0.1, 0.3, 0.4, 0.2])
y_pred_mpc = MPCTensor(secret=y_pred, session=session)
res = mse_loss(y_mpc, y_pred_mpc, reduction)
res_expected = torch.nn.functional.mse_loss(y_secret,
y_pred,
reduction=reduction)
assert np.allclose(res.reconstruct(), res_expected, atol=1e-4)
def test_ops_mpc_public(get_clients, nr_clients, op_str) -> None:
clients = get_clients(nr_clients)
session = Session(parties=clients)
SessionManager.setup_mpc(session)
x_secret = torch.Tensor([[0.125, -1.25], [-4.25, 4]])
if op_str == "truediv":
y_secret = torch.Tensor([[2, 3], [4, 5]]).long()
else:
y_secret = torch.Tensor([[4.5, -2.5], [5, 2.25]])
x = MPCTensor(secret=x_secret, session=session)
op = getattr(operator, op_str)
expected_result = op(x_secret, y_secret)
result = op(x, y_secret).reconstruct()
assert torch.allclose(result, expected_result, atol=10e-4)
def test_przs_rst_ring_size(get_clients) -> None:
clients = get_clients(3)
falcon = Falcon()
session = Session(protocol=falcon, parties=clients)
SessionManager.setup_mpc(session)
for ring_size in RING_SIZE_TO_TYPE.keys():
rst_pt0 = session.session_ptrs[0].przs_generate_random_share(
shape=(1, 2), ring_size=str(ring_size))
share = rst_pt0.get_copy()
assert share.ring_size == ring_size
assert share.shares[0].dtype == RING_SIZE_TO_TYPE[ring_size]
if ring_size == PRIME_NUMBER:
assert torch.max(torch.cat(share.shares)) <= PRIME_NUMBER - 1
assert torch.min(torch.cat(share.shares)) >= 0
def test_reciprocal(method, get_clients) -> None:
clients = get_clients(2)
session_one = Session(parties=clients)
SessionManager.setup_mpc(session_one)
x_secret = torch.Tensor([-2.0, 6.0, 2.0, 3.0, -5.0, -0.5])
x = MPCTensor(secret=x_secret, session=session_one)
x_secret_reciprocal = torch.reciprocal(x_secret)
x_reciprocal = reciprocal(x, method=method)
assert torch.allclose(x_secret_reciprocal,
x_reciprocal.reconstruct(),
atol=1e-1)
with pytest.raises(ValueError):
x_reciprocal = reciprocal(x, method="exp")
def test_ops_divfloat_exception(get_clients, nr_parties) -> None:
# Define the virtual machines that would be use in the computation
parties = get_clients(nr_parties)
# Setup the session for the computation
session = Session(parties=parties)
SessionManager.setup_mpc(session)
x_secret = torch.Tensor([[0.1, -1], [-4, 4]])
y_secret = torch.Tensor([[4.0, -2.5], [5, 2]])
# 3. Share the secret building an MPCTensor
x = MPCTensor(secret=x_secret, session=session)
y = MPCTensor(secret=y_secret, session=session)
with pytest.raises(ValueError):
x / y
def test_max_pool2d_raises_value_error_kernel_gt_input(get_clients) -> None:
clients = get_clients(2)
session = Session(parties=clients)
SessionManager.setup_mpc(session)
secret = torch.Tensor([[[0.23]]])
mpc = MPCTensor(secret=secret, session=session)
with pytest.raises(ValueError):
sympc.module.nn.max_pool2d(mpc, kernel_size=5, stride=1, padding=1)
with pytest.raises(ValueError):
sympc.module.nn.max_pool2d(mpc,
kernel_size=2,
stride=1,
padding=0,
return_indices=True)
def test_session_default_init() -> None:
"""Test correct initialisation of the Sessin class."""
# Test default init
session = Session()
assert session.uuid is None
assert session.parties == []
assert session.trusted_third_party is None
assert session.crypto_store is None
assert session.protocol is not None
assert isinstance(session.config, Config)
assert session.przs_generators == []
assert session.rank == -1
assert session.session_ptrs == []
assert session.tensor_type == get_type_from_ring(2**64)
assert session.ring_size == 2**64
assert session.min_value == -(2**64) // 2
assert session.max_value == (2**64 - 1) // 2
def test_mpc_print(get_clients) -> None:
clients = get_clients(2)
session = Session(parties=clients)
SessionManager.setup_mpc(session)
x_secret = torch.Tensor([5.0])
x = MPCTensor(secret=x_secret, session=session)
expected = f"[MPCTensor]\nShape: {x_secret.shape}\nRequires Grad: False\n\t|"
expected = (
f"{expected} <VirtualMachineClient: P_0 Client> -> ShareTensorPointer\n\t|"
)
expected = f"{expected} <VirtualMachineClient: P_1 Client> -> ShareTensorPointer"
assert expected == x.__str__()
assert x.__str__() == x.__str__()
def test_reconstruct(get_clients) -> None:
clients = get_clients(2)
session = Session(parties=clients)
SessionManager.setup_mpc(session)
a_rand = 3
a = ShareTensor(data=a_rand, config=Config(encoder_precision=0))
MPCTensor.generate_shares(secret=a, nr_parties=2, tensor_type=torch.long)
MPCTensor.generate_shares(
secret=a_rand, nr_parties=2, config=Config(), tensor_type=torch.long
)
x_secret = torch.Tensor([1, -2, 3.0907, -4.870])
x = MPCTensor(secret=x_secret, session=session)
x = x.reconstruct()
assert np.allclose(x_secret, x)
def test_conv_mpc_mpc(get_clients, nr_clients, bias, stride, padding, op_str) -> None:
clients = get_clients(nr_clients)
session = Session(parties=clients)
SessionManager.setup_mpc(session)
input_secret = torch.ones(1, 1, 4, 4)
weight_secret = torch.ones(1, 1, 2, 2)
input = MPCTensor(secret=input_secret, session=session)
weight = MPCTensor(secret=weight_secret, session=session)
kwargs = {"bias": bias, "stride": stride, "padding": padding}
op = getattr(MPCTensor, op_str)
result = op(input, weight, **kwargs).reconstruct()
op = getattr(torch.nn.functional, op_str)
expected_result = op(input_secret, weight_secret, **kwargs)
assert np.allclose(result, expected_result, rtol=10e-4)
def test_mul_private_matrix(get_clients, security, base, precision):
parties = get_clients(3)
protocol = Falcon(security)
config = Config(encoder_base=base, encoder_precision=precision)
session = Session(protocol=protocol, parties=parties, config=config)
SessionManager.setup_mpc(session)
secret1 = torch.tensor([[-100.25, 20.3, 30.12], [-50.1, 100.217, 1.2],
[1032.15, -323.56, 15.15]])
secret2 = torch.tensor([[-1, 0.28, 3], [-9, 10.18, 1], [32, -23, 5]])
tensor1 = MPCTensor(secret=secret1, session=session)
tensor2 = MPCTensor(secret=secret2, session=session)
result = tensor1 * tensor2
expected_res = secret1 * secret2
assert np.allclose(result.reconstruct(), expected_res, atol=1e-3)
def test_max_dim(dim, keepdim, get_clients) -> None:
clients = get_clients(2)
session = Session(parties=clients)
SessionManager.setup_mpc(session)
secret = torch.Tensor([[[1, 2], [3, -1], [4, 5]], [[2, 5], [5, 1], [6, 42]]])
x = MPCTensor(secret=secret, session=session)
max_val, max_idx_val = x.max(dim=dim, keepdim=keepdim)
assert isinstance(x, MPCTensor), "Expected argmax to be MPCTensor"
res_idx = max_idx_val.reconstruct()
res_max = max_val.reconstruct()
expected_max, expected_indices = secret.max(dim=dim, keepdim=keepdim)
assert (
res_idx == expected_indices
).all(), f"Expected indices for maximum to be {expected_indices}"
assert (res_max == expected_max).all(), f"Expected argmax to be {expected_max}"
def test_ops_public_mul_integer_parties(get_clients, parties, security):
config = Config(encoder_base=1, encoder_precision=0)
parties = get_clients(parties)
protocol = Falcon(security)
session = Session(protocol=protocol, parties=parties, config=config)
SessionManager.setup_mpc(session)
secret = torch.tensor([[-100, 20, 30], [-90, 1000, 1], [1032, -323, 15]])
value = 8
op = getattr(operator, "mul")
tensor = MPCTensor(secret=secret, session=session)
shares = [op(share, value) for share in tensor.share_ptrs]
result = MPCTensor(shares=shares, session=session)
assert (result.reconstruct() == (secret * value)).all()
def test_truncation_algorithm1(get_clients, base, precision) -> None:
parties = get_clients(3)
falcon = Falcon("semi-honest")
config = Config(encoder_base=base, encoder_precision=precision)
session = Session(parties=parties, protocol=falcon, config=config)
SessionManager.setup_mpc(session)
x = torch.tensor([[1.24, 4.51, 6.87], [7.87, 1301, 541]])
x_mpc = MPCTensor(secret=x, session=session)
result = ABY3.truncate(x_mpc, session, session.ring_size, session.config)
fp_encoder = FixedPointEncoder(base=session.config.encoder_base,
precision=session.config.encoder_precision)
expected_res = x_mpc.reconstruct(decode=False) // fp_encoder.scale
expected_res = fp_encoder.decode(expected_res)
assert np.allclose(result.reconstruct(), expected_res, atol=1e-3)
def test_bit_decomposition_ttp(get_clients, security_type) -> None:
parties = get_clients(3)
falcon = Falcon(security_type=security_type)
session = Session(parties=parties, protocol=falcon)
SessionManager.setup_mpc(session)
secret = torch.tensor([[-1, 12], [-32, 45], [98, -5624]])
x = MPCTensor(secret=secret, session=session)
b_sh = ABY3.bit_decomposition_ttp(x, session)
ring_size = x.session.ring_size
tensor_type = x.session.tensor_type
ring_bits = get_nr_bits(ring_size)
result = torch.zeros(size=x.shape, dtype=tensor_type)
for i in range(ring_bits):
result |= b_sh[i].reconstruct(decode=False).type(tensor_type) << i
exp_res = torch.tensor([[-65536, 786432], [-2097152, 2949120],
[6422528, -368574464]])
assert (result == exp_res).all()
def test_send_get(get_clients, precision=12, base=4) -> None:
client = get_clients(1)[0]
protocol = Falcon("semi-honest")
session = Session(protocol=protocol, parties=[client])
SessionManager.setup_mpc(session)
share1 = torch.Tensor([1.4, 2.34, 3.43])
share2 = torch.Tensor([1, 2, 3])
share3 = torch.Tensor([1.4, 2.34, 3.43])
session_uuid = session.rank_to_uuid[0]
x_share = ReplicatedSharedTensor(shares=[share1, share2, share3],
session_uuid=session_uuid)
x_ptr = x_share.send(client)
result = x_ptr.get()
assert result == x_share
def test_session_custom_init() -> None:
config = Config()
session = Session(parties=["alice", "bob"],
ring_size=2**32,
config=config,
ttp="TTP")
assert session.uuid is None
assert session.parties == ["alice", "bob"]
assert session.trusted_third_party == "TTP"
assert session.crypto_store is None
assert session.protocol is not None
assert session.config == config
assert session.przs_generators == []
assert session.rank == -1
assert session.session_ptrs == []
assert session.tensor_type == get_type_from_ring(2**32)
assert session.ring_size == 2**32
assert session.min_value == -(2**32) // 2
assert session.max_value == (2**32 - 1) // 2
def test_primitive_logging_beaver_matmul(get_clients) -> None:
clients = get_clients(2)
session = Session(parties=clients)
SessionManager.setup_mpc(session)
p_kwargs = {"a_shape": (2, 3), "b_shape": (3, 10)}
g_kwargs = {"a_shape": (2, 3), "b_shape": (3, 10), "nr_parties": 2}
CryptoPrimitiveProvider.start_logging()
CryptoPrimitiveProvider.generate_primitives(
session=session,
op_str="beaver_matmul",
p_kwargs=p_kwargs,
g_kwargs=g_kwargs,
)
primitive_log = CryptoPrimitiveProvider.stop_logging()
expected_log = {"beaver_matmul": [(p_kwargs, g_kwargs)]}
assert expected_log == primitive_log
def test_generate_primitive_from_dict_beaver_conv2d(get_clients) -> None:
clients = get_clients(2)
session = Session(parties=clients)
SessionManager.setup_mpc(session)
primitive_log = {
"beaver_conv2d": [(
{
"a_shape": (1, 1, 28, 28),
"b_shape": (5, 1, 5, 5)
},
{
"session": session,
"a_shape": (1, 1, 28, 28),
"b_shape": (5, 1, 5, 5),
"nr_parties": 2,
},
)]
}
CryptoPrimitiveProvider.generate_primitive_from_dict(
primitive_log=primitive_log, session=session)
a_shape = (1, 1, 28, 28)
b_shape = (5, 1, 5, 5)
key = f"beaver_conv2d_{a_shape}_{b_shape}"
store_client_1 = session.session_ptrs[0].crypto_store.store.get()
store_client_2 = session.session_ptrs[1].crypto_store.store.get()
a_shape_client_1 = tuple(store_client_1.get(key)[0][0].shape)
b_shape_client_1 = tuple(store_client_1.get(key)[0][1].shape)
assert a_shape == a_shape_client_1
assert b_shape == b_shape_client_1
a_shape_client_2 = tuple(store_client_2.get(key)[0][0].shape)
b_shape_client_2 = tuple(store_client_2.get(key)[0][1].shape)
assert a_shape == a_shape_client_2
assert b_shape == b_shape_client_2
def test_generate_and_transfer_primitive(
get_clients: Callable, nr_parties: int, nr_instances: int
) -> None:
parties = get_clients(nr_parties)
session = Session(parties=parties)
SessionManager.setup_mpc(session)
g_kwargs = {"nr_parties": nr_parties}
CryptoPrimitiveProvider.generate_primitives(
"test",
n_instances=nr_instances,
sessions=session.session_ptrs,
g_kwargs=g_kwargs,
p_kwargs={},
)
for i in range(nr_parties):
remote_crypto_store = session.session_ptrs[i].crypto_store
primitives = remote_crypto_store.get_primitives_from_store("test").get()
assert primitives == [tuple(i for _ in range(PRIMITIVE_NR_ELEMS))]
def test_generate_primitive(get_clients: Callable, nr_parties: int,
nr_instances: int) -> None:
parties = get_clients(nr_parties)
session = Session(parties=parties)
SessionManager.setup_mpc(session)
g_kwargs = {"nr_parties": nr_parties, "nr_instances": nr_instances}
res = CryptoPrimitiveProvider.generate_primitives(
"test",
sessions=session.session_ptrs,
g_kwargs=g_kwargs,
p_kwargs=None,
)
assert isinstance(res, list)
assert len(res) == nr_parties
for i, primitives in enumerate(res):
for primitive in primitives:
assert primitive == tuple(i for _ in range(PRIMITIVE_NR_ELEMS))
