# Buggy code
def private_compare(x_bit_sh, r, beta, L):
    """
    Perform privately x > r

    args:
        x (AdditiveSharedTensor): the private tensor
        r (MultiPointerTensor): the threshold commonly held by alice and bob
        beta (MultiPointerTensor): a boolean commonly held by alice and bob to
            hide the result of computation for the crypto provider
        L(int): field size for r

    return:
        β′ = β ⊕ (x > r).
    """
    assert isinstance(x_bit_sh, sy.AdditiveSharingTensor)
    assert isinstance(r, sy.MultiPointerTensor)
    assert isinstance(beta, sy.MultiPointerTensor)
    # Would it be safer to have a different r/beta for each value in the tensor?

    alice, bob = x_bit_sh.locations
    crypto_provider = x_bit_sh.crypto_provider
    p = x_bit_sh.field

    # the commented out numbers below correspond to the
    # line numbers in Algorithm 3 of the SecureNN paper
    # https://eprint.iacr.org/2018/442.pdf

    # Common randomess
    s = torch.randint(1, p, x_bit_sh.shape).send(alice, bob, **no_wrap)
    u = torch.randint(1, p, x_bit_sh.shape).send(alice, bob, **no_wrap)
    perm = torch.randperm(x_bit_sh.shape[-1]).send(alice, bob, **no_wrap)

    j = sy.MultiPointerTensor(
        children=[
            torch.tensor([0]).send(alice, **no_wrap),
            torch.tensor([1]).send(bob, **no_wrap),
        ]
    )

    # 1)
    t = r + 1
    t_bit = decompose(t, L)
    r_bit = decompose(r, L)

    # if beta == 0
    # 5)
    w = x_bit_sh + (j * r_bit) - (2 * r_bit * x_bit_sh)
    # 6)
    wc = w.flip(-1).cumsum(-1).flip(-1) - w
    c_beta0 = -x_bit_sh + (j * r_bit) + j + wc

    # elif beta == 1 AND r != 2^l- 1
    # 8)
    w = x_bit_sh + (j * t_bit) - (2 * t_bit * x_bit_sh)
    # 9)
    wc = w.flip(-1).cumsum(-1).flip(-1) - w
    c_beta1 = x_bit_sh + (-j * t_bit) + j + wc

    # else
    # 11)
    c_igt1 = (1 - j) * (u + 1) - (j * u)
    c_ie1 = (1 - 2 * j) * u

    l1_mask = torch.zeros(x_bit_sh.shape).long()
    l1_mask[..., 0] = 1
    l1_mask = l1_mask.send(alice, bob, **no_wrap)
    # c_else = if i == 1 c_ie1 else c_igt1
    c_else = (l1_mask * c_ie1) + ((1 - l1_mask) * c_igt1)

    # Mask for the case r == 2^l −1
    r_mask = (r == get_r_mask(L)).long()
    r_mask = r_mask.unsqueeze(-1)

    # Mask combination to execute the if / else statements of 4), 7), 10)
    c = (
        (1 - beta) * c_beta0
        + (beta * (1 - r_mask)) * c_beta1
        + (beta * r_mask) * c_else
    )

    # 14)
    # Hide c values
    mask = s * c

    # Permute the mask
    # I have to create idx because Ellipsis are still not supported
    # (I would like to do permuted_mask = mask[..., perm])
    idx = [slice(None)] * (len(x_bit_sh.shape) - 1) + [perm]
    permuted_mask = mask[idx]
    # Send it to another worker
    # We do this because we can't allow the local worker to get and see permuted_mask
    # because otherwise it can inverse the permutation and remove s to get c.
    # So opening the permuted_mask should be made by a worker which doesn't have access to the randomness
    remote_mask = permuted_mask.wrap().send(crypto_provider, **no_wrap)

    # 15)
    d_ptr = remote_mask.remote_get()
    beta_prime = (d_ptr == 0).sum(-1)

    # Get result back
    res = beta_prime.get()
    return res


# Traceback (TypeError)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-7-902fb9b77716> in <module>
12 x = th.tensor([1.1,2.0,3.2,4.0]).fix_precision().share(*workers, crypto_provider=sw, requires_grad=True)
13 m=th.nn.Linear(4, 1).fix_precision().share(*workers, crypto_provider=sw, requires_grad=True)
---> 14 o=m(x)

/usr/local/lib/python3.7/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
491             result = self._slow_forward(*input, **kwargs)
492         else:
--> 493             result = self.forward(*input, **kwargs)
494         for hook in self._forward_hooks.values():
495             hook_result = hook(self, input, result)

/usr/local/lib/python3.7/site-packages/torch/nn/modules/linear.py in forward(self, input)
90     @weak_script_method
91     def forward(self, input):
---> 92         return F.linear(input, self.weight, self.bias)
93
94     def extra_repr(self):

/usr/local/lib/python3.7/site-packages/syft/frameworks/torch/hook/hook.py in overloaded_func(*args, **kwargs)
745                 handle_func_command = TorchTensor.handle_func_command
746
--> 747             response = handle_func_command(command)
748
749             return response

/usr/local/lib/python3.7/site-packages/syft/frameworks/torch/tensors/interpreters/native.py in handle_func_command(cls, command)
310             new_command = (cmd, None, new_args, new_kwargs)
311             # Send it to the appropriate class and get the response
--> 312             response = new_type.handle_func_command(new_command)
313             # Put back the wrappers where needed
314             response = syft.frameworks.torch.hook_args.hook_response(

/usr/local/lib/python3.7/site-packages/syft/frameworks/torch/tensors/interpreters/autograd.py in handle_func_command(cls, command)
236             # Try to get recursively the attributes in cmd = "<attr1>.<attr2>.<attr3>..."
237             cmd = cls.rgetattr(cls, cmd)
--> 238             return cmd(*args, **kwargs)
239         except AttributeError:
240             pass

/usr/local/lib/python3.7/site-packages/syft/frameworks/torch/tensors/interpreters/autograd.py in linear(*args)
201                     Un-hook the function to have its detailed behaviour
202                     """
--> 203                     return torch.nn.functional.native_linear(*args)
204
205                 module.linear = linear

/usr/local/lib/python3.7/site-packages/torch/nn/functional.py in linear(input, weight, bias)
1406         ret = torch.addmm(bias, input, weight.t())
1407     else:
-> 1408         output = input.matmul(weight.t())
1409         if bias is not None:
1410             output += bias

/usr/local/lib/python3.7/site-packages/syft/frameworks/torch/tensors/interpreters/autograd.py in method_with_grad(*args, **kwargs)
138                 )
139
--> 140                 result = getattr(new_self, name)(*new_args, **new_kwargs)
141
142                 # Put back SyftTensor on the tensors found in the response

/usr/local/lib/python3.7/site-packages/syft/frameworks/torch/tensors/interpreters/precision.py in matmul(self, *args, **kwargs)
414
415         # Send it to the appropriate class and get the response
--> 416         response = getattr(new_self, "matmul")(*new_args, **new_kwargs)
417
418         # Put back SyftTensor on the tensors found in the response

/usr/local/lib/python3.7/site-packages/syft/frameworks/torch/tensors/interpreters/additive_shared.py in matmul(self, other)
515             return self._public_mul(other, "matmul")
516
--> 517         return self._private_mul(other, "matmul")
518
519     def mm(self, *args, **kwargs):

/usr/local/lib/python3.7/site-packages/syft/frameworks/torch/tensors/interpreters/additive_shared.py in _private_mul(self, other, equation)
413             raise AttributeError("For multiplication a crypto_provider must be passed.")
414
--> 415         shares = spdz.spdz_mul(cmd, self, other, self.crypto_provider, self.field)
416
417         return shares

/usr/local/lib/python3.7/site-packages/syft/frameworks/torch/crypto/spdz.py in spdz_mul(cmd, x_sh, y_sh, crypto_provider, field)
43         j = sy.MultiPointerTensor(children=[j1, j0])
44     else:
---> 45         j = sy.MultiPointerTensor(children=[j1] + j0.child.values())
46
47     delta_b = cmd(delta, b)

TypeError: can only concatenate list (not "dict_values") to list

# Fixed code
def private_compare(x_bit_sh, r, beta, L):
    """
    Perform privately x > r

    args:
        x (AdditiveSharedTensor): the private tensor
        r (MultiPointerTensor): the threshold commonly held by the workers
        beta (MultiPointerTensor): a boolean commonly held by the workers to
            hide the result of computation for the crypto provider
        L(int): field size for r

    return:
        β′ = β ⊕ (x > r).
    """
    assert isinstance(x_bit_sh, sy.AdditiveSharingTensor)
    assert isinstance(r, sy.MultiPointerTensor)
    assert isinstance(beta, sy.MultiPointerTensor)
    # Would it be safer to have a different r/beta for each value in the tensor?

    workers = x_bit_sh.locations
    crypto_provider = x_bit_sh.crypto_provider
    p = x_bit_sh.field

    # the commented out numbers below correspond to the
    # line numbers in Algorithm 3 of the SecureNN paper
    # https://eprint.iacr.org/2018/442.pdf

    # Common randomess
    s = torch.randint(1, p, x_bit_sh.shape).send(*workers, **no_wrap)
    u = torch.randint(1, p, x_bit_sh.shape).send(*workers, **no_wrap)
    perm = torch.randperm(x_bit_sh.shape[-1]).send(*workers, **no_wrap)

    j = sy.MultiPointerTensor(
        children=[
            torch.tensor([int(i == 0)]).send(w, **no_wrap)
            for i, w in enumerate(workers)
        ]
    )

    # 1)
    t = r + 1
    t_bit = decompose(t, L)
    r_bit = decompose(r, L)

    # if beta == 0
    # 5)
    w = x_bit_sh + (j * r_bit) - (2 * r_bit * x_bit_sh)
    # 6)
    wc = w.flip(-1).cumsum(-1).flip(-1) - w
    c_beta0 = -x_bit_sh + (j * r_bit) + j + wc

    # elif beta == 1 AND r != 2^l- 1
    # 8)
    w = x_bit_sh + (j * t_bit) - (2 * t_bit * x_bit_sh)
    # 9)
    wc = w.flip(-1).cumsum(-1).flip(-1) - w
    c_beta1 = x_bit_sh + (-j * t_bit) + j + wc

    # else
    # 11)
    c_igt1 = (1 - j) * (u + 1) - (j * u)
    c_ie1 = (1 - 2 * j) * u

    l1_mask = torch.zeros(x_bit_sh.shape).long()
    l1_mask[..., 0] = 1
    l1_mask = l1_mask.send(*workers, **no_wrap)
    # c_else = if i == 1 c_ie1 else c_igt1
    c_else = (l1_mask * c_ie1) + ((1 - l1_mask) * c_igt1)

    # Mask for the case r == 2^l −1
    r_mask = (r == get_r_mask(L)).long()
    r_mask = r_mask.unsqueeze(-1)

    # Mask combination to execute the if / else statements of 4), 7), 10)
    c = (
        (1 - beta) * c_beta0
        + (beta * (1 - r_mask)) * c_beta1
        + (beta * r_mask) * c_else
    )

    # 14)
    # Hide c values
    mask = s * c

    # Permute the mask
    # I have to create idx because Ellipsis are still not supported
    # (I would like to do permuted_mask = mask[..., perm])
    idx = [slice(None)] * (len(x_bit_sh.shape) - 1) + [perm]
    permuted_mask = mask[idx]
    # Send it to another worker
    # We do this because we can't allow the local worker to get and see permuted_mask
    # because otherwise it can inverse the permutation and remove s to get c.
    # So opening the permuted_mask should be made by a worker which doesn't have access to the randomness
    remote_mask = permuted_mask.wrap().send(crypto_provider, **no_wrap)

    # 15)
    d_ptr = remote_mask.remote_get()
    beta_prime = (d_ptr == 0).sum(-1)

    # Get result back
    res = beta_prime.get()
    return res