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