def test_max_multiple_max(get_clients) -> None:
clients = get_clients(2)
session = Session(parties=clients)
SessionManager.setup_mpc(session)
x = MPCTensor(secret=torch.Tensor([1, 2, 3, -1, 3]), session=session)
with pytest.raises(ValueError):
x.argmax()
def test_max(get_clients) -> None:
clients = get_clients(2)
session = Session(parties=clients)
SessionManager.setup_mpc(session)
secret = torch.Tensor([1, 2, 3, -1, -3])
x = MPCTensor(secret=secret, session=session)
max_val = x.max()
assert isinstance(x, MPCTensor), "Expected argmax to be MPCTensor"
expected = secret.max()
res = max_val.reconstruct()
assert res == expected, f"Expected argmax to be {expected}"
def test_cat(get_clients):
clients = get_clients(2)
x_secret = torch.Tensor([0.0, 1, -2, 3, -4])
y_secret = torch.Tensor([-4, 3, -2, 1, 0.0])
secret_concatenated = torch.cat([x_secret, y_secret])
session = Session(parties=clients)
SessionManager.setup_mpc(session)
x = MPCTensor(secret=x_secret, session=session)
y = MPCTensor(secret=y_secret, session=session)
concatenated = cat([x, y])
assert (secret_concatenated == concatenated.reconstruct()).all()
def test_argmax_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)
argmax_val = x.argmax(dim=dim, keepdim=keepdim)
assert isinstance(x, MPCTensor), "Expected argmax to be MPCTensor"
res = argmax_val.reconstruct()
expected = secret.argmax(dim=dim, keepdim=keepdim).float()
assert (res == expected).all(), f"Expected argmax to be {expected}"
def stack(tensors: List, dim: int = 0) -> MPCTensor:
"""Concatenates a sequence of tensors along a new dimension.
Args:
tensors (List): sequence of tensors to stacks
dim (int): dimension to insert. Has to be between 0 and the number of
dimensions of concatenated tensors (inclusive)
Returns:
MPCTensor: calculated MPCTensor
"""
session = tensors[0].session
args = list(
zip([str(uuid) for uuid in session.rank_to_uuid.values()],
*[tensor.share_ptrs for tensor in tensors]))
stack_shares = parallel_execution(stack_share_tensor,
session.parties)(args)
from sympc.tensor import MPCTensor
expected_shape = torch.stack(
[torch.empty(each_tensor.shape) for each_tensor in tensors],
dim=dim).shape
result = MPCTensor(shares=stack_shares,
session=session,
shape=expected_shape)
return result
def cat(tensors: List, dim: int = 0) -> MPCTensor:
"""Concatenates the given sequence of seq tensors in the given dimension.
Args:
tensors (List): sequence of tensors to concatenate
dim (int): the dimension over which the tensors are concatenated
Returns:
MPCTensor: calculated MPCTensor
"""
session = tensors[0].session
args = list(
zip([str(uuid) for uuid in session.rank_to_uuid.values()],
*[tensor.share_ptrs for tensor in tensors]))
stack_shares = parallel_execution(cat_share_tensor, session.parties)(args)
from sympc.tensor import MPCTensor
expected_shape = torch.cat(
[torch.empty(each_tensor.shape) for each_tensor in tensors],
dim=dim).shape
result = MPCTensor(shares=stack_shares,
session=session,
shape=expected_shape)
return result
def backward(ctx: Dict[str, Any], grad: MPCTensor) -> Tuple[MPCTensor]:
"""Perform the backward pass for the matrix multiplication operation.
Args:
ctx (Dict[str, Any]): Context used to retrieve the information for the backward pass
grad (MPCTensor): The gradient that came from the child nodes
Returns:
(x_grad, y_grad) (Tuple[MPCTensor]): The gradients passed to the X and Y nodes.
Raises:
ValueError: if gradient shape does not match X and Y shape
"""
x, y = ctx["x"], ctx["y"]
x_grad = grad.clone()
y_grad = grad.clone()
if len(x.shape) < 2:
x = x.unsqueeze(0)
x_grad = x_grad.unsqueeze(0)
if len(y.shape) < 2:
y = y.unsqueeze(1)
y_grad = y_grad.unsqueeze(1)
x_grad = x_grad @ y.t()
y_grad = x.t() @ y_grad
if x.shape != x_grad.shape or y.shape != y_grad.shape:
raise ValueError(
"The gradient shape and the shape of X and Y should be the same!"
)
return x_grad, y_grad
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 forward(ctx: Dict[str, Any], x: MPCTensor) -> MPCTensor:
"""Perform the feedforward and compute the result for the transpose operation.
Args:
ctx (Dict[str, Any]): Context used to save information needed in the backward pass
x (MPCTensor): The operand to apply the transpose on
Returns:
x.t() (MPCTensor): The result after applying transpose
"""
return x.t()
def test_exception_value_error(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)
with pytest.raises(ValueError):
reciprocal(x, method="exp")
def backward(ctx: Dict[str, Any], grad: MPCTensor) -> MPCTensor:
"""Perform the backward pass for the transpose operation.
Args:
ctx (Dict[str, Any]): Context used to retrieve the information for the backward pass
grad (MPCTensor): The gradient that came from the child nodes
Returns:
res_grad (MPCTensor): The gradients passed to the parent node
"""
res_grad = grad.t()
return res_grad
def test_softmax(get_clients, dim) -> None:
clients = get_clients(2)
x_secret = torch.arange(-6, 6, dtype=torch.float).view(3, 4)
x_secret_softmax = F.softmax(x_secret, dim=dim)
session = Session(parties=clients)
SessionManager.setup_mpc(session)
x = MPCTensor(secret=x_secret, session=session)
x_softmax = softmax(x, dim=dim)
assert torch.allclose(x_secret_softmax, x_softmax.reconstruct(), atol=1e-2)
def test_exception_value_error(get_clients) -> None:
clients = get_clients(2)
x_secret = torch.Tensor([0.0, 1, -2, 3, -4])
torch.tanh(x_secret)
session = Session(parties=clients)
SessionManager.setup_mpc(session)
x = MPCTensor(secret=x_secret, session=session)
with pytest.raises(ValueError):
tanh(x, method="exp")
def test_softmax_single_along_dim(get_clients) -> None:
clients = get_clients(2)
x_secret = torch.arange(4, dtype=torch.float).view(4, 1)
x_secret_softmax = F.softmax(x_secret)
session = Session(parties=clients)
SessionManager.setup_mpc(session)
x = MPCTensor(secret=x_secret, session=session)
x_softmax = softmax(x)
assert torch.allclose(x_secret_softmax, x_softmax.reconstruct(), atol=1e-2)
def test_sigmoid(get_clients, method) -> None:
clients = get_clients(2)
x_secret = torch.Tensor([0.0, 1, -2, 3, -4])
x_secret_sigmoid = torch.sigmoid(x_secret)
session = Session(parties=clients)
SessionManager.setup_mpc(session)
x = MPCTensor(secret=x_secret, session=session)
x_sigmoid = sigmoid(x, method)
assert torch.allclose(x_secret_sigmoid, x_sigmoid.reconstruct(), atol=1e-1)
def backward(ctx: Dict[str, Any], grad: MPCTensor) -> MPCTensor:
"""Perform the backward pass for the reshape operation.
Args:
ctx (Dict[str, Any]): Context used to retrieve the information for the backward pass
grad (MPCTensor): The gradient that came from the child nodes
Returns:
grad (MPCTensor): The gradients passed to the X node.
"""
shape = tuple(ctx["x_shape"])
return grad.reshape(shape)
def forward(ctx: Dict[str, Any], x: MPCTensor) -> MPCTensor:
"""Perform the feedforward and compute the result for the sigmoid operation.
Args:
ctx (Dict[str, Any]): Context used to save information needed in the backward pass
x (MPCTensor): The operand on which to apply the sigmoid function
Returns:
sigmoid(x) (MPCTensor): The sigmoid approximation applied on the input
"""
grad = x.sigmoid()
ctx["probabilities"] = grad
return grad
def forward(ctx: Dict[str, Any], x: MPCTensor) -> MPCTensor:
"""Perform the feedforward and compute the result for the summation operation.
Args:
ctx (Dict[str, Any]): Context used to save information needed in the backward pass
x (MPCTensor): The operand on which to apply the sum function
Returns:
sum(x) (MPCTensor): The summation of all the elements from the input
"""
ctx["x_shape"] = x.shape
total_sum = x.sum()
return total_sum
def forward(ctx: Dict[str, Any], x: MPCTensor) -> MPCTensor:
"""Perform the feedforward and compute the result for the ReLU operation.
Args:
ctx (Dict[str, Any]): Context used to save information needed in the backward pass
x (MPCTensor): The operand on which to apply the sigmoid function
Returns:
relu(x) (MPCTensor): The sigmoid approximation applied on the input
"""
mask = APPROXIMATIONS["sign"](x.gt(0.0))
ctx["mask"] = mask
return x * mask
def forward(ctx: Dict[str, Any], x: MPCTensor, shape: tuple) -> MPCTensor:
"""Perform the feedforward and compute the result for the reshape operation.
Args:
ctx (Dict[str, Any]): Context used to save information needed in the backward pass
x (MPCTensor): the MPCTensor to be reshaped
shape (tuple): the new shape
Returns:
res (MPCTensor): The result of the reshape operation
"""
ctx["x_shape"] = x.shape
return x.reshape(shape)
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)
def test_log(get_clients) -> None:
clients = get_clients(2)
x_secret = torch.Tensor([0.1, 0.5, 2, 5, 10])
x_secret_log = torch.log(x_secret)
# with custom precision
config = Config(encoder_precision=20)
session = Session(parties=clients, config=config)
SessionManager.setup_mpc(session)
x = MPCTensor(secret=x_secret, session=session)
x_log = log(x)
assert torch.allclose(x_secret_log, x_log.reconstruct(), atol=1e-1)
def test_tanh(get_clients) -> None:
clients = get_clients(2)
x_secret = torch.Tensor([0.0, 1, -2, 3, -4])
x_secret_tanh = torch.tanh(x_secret)
session = Session(parties=clients)
SessionManager.setup_mpc(session)
x = MPCTensor(secret=x_secret, session=session)
x_tanh = tanh(x, method="sigmoid")
assert torch.allclose(x_secret_tanh, x_tanh.reconstruct(), atol=1e-2)
with pytest.raises(ValueError):
x_tanh = tanh(x, method="exp")
def forward(
ctx: Dict[str, Any], x: MPCTensor, start: int = 0, end: int = -1
) -> MPCTensor:
"""Perform the feedforward and compute the result for the multiplication operation.
Args:
ctx (Dict[str, Any]): Context used to save information needed in the backward pass
x (MPCTensor): 1st operand for the multiplication operation
start (int): Start dimension for the flatten operation
end (int): Final dimension for the flatten operation
Returns:
res (MPCTensor): The result of the flatten operation
"""
ctx["x_shape"] = x.shape
return x.flatten(start_dim=start, end_dim=end)
def backward(ctx: Dict[str, Any], grad: MPCTensor) -> Tuple[MPCTensor]:
"""Perform the backward pass for the conv2d operation.
Args:
ctx (Dict[str, Any]): Context used to retrieve the information for the backward pass
grad (MPCTensor): The gradient that came from the child nodes
Returns:
(input_grad, weight_grad) (Tuple[MPCTensor]): The gradients passed
to the input and kernal nodes.
"""
input = ctx["input"]
weight = ctx["weight"]
stride = ctx["stride"]
padding = ctx["padding"]
dilation = ctx["dilation"]
groups = ctx["groups"]
weight_size = (weight.shape[2], weight.shape[3])
in_channels = input.shape[1]
out_channels = grad.shape[1]
min_batch = input.shape[0]
# Gradient w.r.t input of the Conv.
common_args = [
tuple(input.shape),
stride,
padding,
weight_size,
dilation,
grad.session,
]
args = [[el] + common_args for el in grad.share_ptrs]
shares = parallel_execution(
GradConv2d.get_grad_input_padding, grad.session.parties
)(args)
grad_input_padding = MPCTensor(shares=shares, session=grad.session)
output_padding_tensor = grad_input_padding.reconstruct()
output_padding_tensor /= grad.session.nr_parties
output_padding = tuple(output_padding_tensor.to(torch.int).tolist())
input_grad = grad.conv_transpose2d(
weight, None, stride, output_padding, dilation, groups
)
# Gradient w.r.t weights of the Conv.
grad = grad.repeat(1, in_channels // groups, 1, 1)
grad = grad.view(grad.shape[0] * grad.shape[1], 1, grad.shape[2], grad.shape[3])
input = input.view(
1, input.shape[0] * input.shape[1], input.shape[2], input.shape[3]
)
weight_grad = input.conv2d(
weight=grad,
bias=None,
dilation=stride,
padding=padding,
stride=dilation,
groups=in_channels * min_batch,
)
weight_grad = weight_grad.view(
min_batch,
weight_grad.shape[1] // min_batch,
weight_grad.shape[2],
weight_grad.shape[3],
)
weight_grad = (
weight_grad.sum(0)
.view(
in_channels // groups,
out_channels,
weight_grad.shape[2],
weight_grad.shape[3],
)
.transpose(0, 1)
)
weight_grad = weight_grad.narrow(2, 0, weight_size[1])
weight_grad = weight_grad.narrow(3, 0, weight_size[0])
return input_grad, weight_grad
def helper_argmax(
x: MPCTensor,
dim: Optional[Union[int, Tuple[int]]] = None,
keepdim: bool = False,
one_hot: bool = False,
) -> MPCTensor:
"""Compute argmax using pairwise comparisons. Makes the number of rounds fixed, here it is 2.
This is inspired from CrypTen.
Args:
x (MPCTensor): the MPCTensor on which to compute helper_argmax on
dim (Union[int, Tuple[int]): compute argmax over a specific dimension(s)
keepdim (bool): when one_hot is true, keep all the dimensions of the tensor
one_hot (bool): return the argmax as a one hot vector
Returns:
Given the args, it returns a one hot encoding (as an MPCTensor) or the index
of the maximum value
Raises:
ValueError: In case more max values are found and we need to return the index
"""
# for each share in MPCTensor
# do the algorithm portrayed in paper (helper_argmax_pairwise)
# results in creating two matrices and subtraction them
session = x.session
prep_x = x.flatten() if dim is None else x
args = [[str(uuid), share_ptr_tensor,
dim] for uuid, share_ptr_tensor in zip(
session.rank_to_uuid.values(), prep_x.share_ptrs)]
shares = parallel_execution(helper_argmax_pairwise, session.parties)(args)
res_shape = shares[0].shape.get()
x_pairwise = MPCTensor(shares=shares, session=x.session, shape=res_shape)
# with the MPCTensor tensor we check what entries are positive
# then we check what columns of M matrix have m-1 non-zero entries after comparison
# (by summing over cols)
pairwise_comparisons = x_pairwise >= 0
# re-compute row_length
_dim = -1 if dim is None else dim
row_length = x.shape[_dim] if x.shape[_dim] > 1 else 2
result = pairwise_comparisons.sum(0)
result = result >= (row_length - 1)
res_shape = res_shape[1:] # Remove the leading dimension because of sum(0)
if not one_hot:
if dim is None:
check = result * torch.Tensor(
[i for i in range(np.prod(res_shape))])
else:
size = [1 for _ in range(len(res_shape))]
size[dim] = res_shape[dim]
check = result * torch.Tensor([i for i in range(res_shape[_dim])
]).view(size)
if dim is not None:
argmax = check.sum(dim=dim, keepdim=keepdim)
else:
argmax = check.sum()
if (argmax >= row_length).reconstruct():
# In case we have 2 max values, rather then returning an invalid index
# we raise an exception
raise ValueError("There are multiple argmax values")
result = argmax
return result
def backward(ctx: Dict[str, Any], grad: MPCTensor) -> Tuple[MPCTensor]:
"""Perform the backward pass for the conv2d operation.
Args:
ctx (Dict[str, Any]): Context used to retrieve the information for the backward pass
grad (MPCTensor): The gradient that came from the child nodes
Returns:
(input_grad, weight_grad) (Tuple[MPCTensor, MPCTensor]): The gradients passed
to the input and kernal nodes.
"""
x = ctx["x"]
weight = ctx["weight"]
stride = ctx["stride"]
padding = ctx["padding"]
dilation = ctx["dilation"]
groups = ctx["groups"]
weight_size = (weight.shape[2], weight.shape[3])
in_channels = x.shape[1]
out_channels = grad.shape[1]
min_batch = x.shape[0]
output_padding = torch.nn.grad._grad_input_padding(
grad_output=torch.empty(grad.shape),
input_size=x.shape,
stride=(stride, stride),
padding=(padding, padding),
kernel_size=weight_size,
dilation=(dilation, dilation),
)
input_grad = grad.conv_transpose2d(weight, None, stride,
output_padding, dilation, groups)
# Gradient w.r.t weights of the Conv.
grad = grad.repeat(1, in_channels // groups, 1, 1)
grad = grad.view(grad.shape[0] * grad.shape[1], 1, grad.shape[2],
grad.shape[3])
x = x.view(1, x.shape[0] * x.shape[1], x.shape[2], x.shape[3])
weight_grad = x.conv2d(
weight=grad,
bias=None,
dilation=stride,
padding=padding,
stride=dilation,
groups=in_channels * min_batch,
)
weight_grad = weight_grad.view(
min_batch,
weight_grad.shape[1] // min_batch,
weight_grad.shape[2],
weight_grad.shape[3],
)
weight_grad = (weight_grad.sum(0).view(
in_channels // groups,
out_channels,
weight_grad.shape[2],
weight_grad.shape[3],
).transpose(0, 1))
weight_grad = weight_grad.narrow(2, 0, weight_size[1])
weight_grad = weight_grad.narrow(3, 0, weight_size[0])
return input_grad, weight_grad
