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))