# Buggy code
def maxpool2d(a_sh, kernel_size: int = 1, stride: int = 1, padding: int = 0):
    """Applies a 2D max pooling over an input signal composed of several input planes.
    This interface is similar to torch.nn.MaxPool2D.
    Args:
        kernel_size: the size of the window to take a max over
        stride: the stride of the window
        padding: implicit zero padding to be added on both sides
    """
    assert a_sh.dtype != "custom", (
        "`custom` dtype shares are unsupported in SecureNN, use dtype = `long` or `int` instead"
    )
    assert len(a_sh.shape) == 4

    # Change to tuple if not one
    kernel = torch.nn.modules.utils._pair(kernel_size)
    stride = torch.nn.modules.utils._pair(stride)
    padding = torch.nn.modules.utils._pair(padding)

    # TODO: support dilation.
    dilation = torch.nn.modules.utils._pair(1)

    # Extract a few useful values
    batch_size, nb_channels, nb_rows_in, nb_cols_in = a_sh.shape

    # Calculate output shapes
    nb_rows_out = int(
        (nb_rows_in + 2 * padding[0] - dilation[0] * (kernel[0] - 1) - 1) / stride[0]
        + 1
    )
    nb_cols_out = int(
        (nb_cols_in + 2 * padding[1] - dilation[1] * (kernel[1] - 1) - 1) / stride[1]
        + 1
    )

    # Apply padding to the input
    if padding != (0, 0):
        a_sh = torch.nn.functional.pad(
            a_sh, (padding[1], padding[1], padding[0], padding[0]), "constant"
        )
        # Update shape after padding
        nb_rows_in += 2 * padding[0]
        nb_cols_in += 2 * padding[1]

    res = []
    # TODO: make this operation more efficient in order to be used with cnn modules.
    for batch in range(batch_size):
        for channel in range(nb_channels):
            for r_in in range(0, nb_rows_in - (kernel[0] - 1), stride[0]):
                for c_in in range(0, nb_cols_in - (kernel[1] - 1), stride[1]):
                    m, _ = maxpool(
                        a_sh[
                            batch,
                            channel,
                            r_in : r_in + kernel[0],
                            c_in : c_in + kernel[1],
                        ].child
                    )
                    res.append(m.wrap())

    res = torch.stack(res).reshape(batch_size, nb_channels, nb_rows_out, nb_cols_out)
    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 maxpool2d(a_sh, kernel_size: int = 1, stride: int = 1, padding: int = 0):
    """Applies a 2D max pooling over an input signal composed of several input planes.
    This interface is similar to torch.nn.MaxPool2D.
    Args:
        kernel_size: the size of the window to take a max over
        stride: the stride of the window
        padding: implicit zero padding to be added on both sides
    """
    assert a_sh.dtype != "custom", (
        "`custom` dtype shares are unsupported in SecureNN, use dtype = `long` or `int` instead"
    )
    assert len(a_sh.shape) == 4

    # Change to tuple if not one
    kernel = torch.nn.modules.utils._pair(kernel_size)
    stride = torch.nn.modules.utils._pair(stride)
    padding = torch.nn.modules.utils._pair(padding)

    # TODO: support dilation.
    dilation = torch.nn.modules.utils._pair(1)

    # Extract a few useful values
    batch_size, nb_channels, nb_rows_in, nb_cols_in = a_sh.shape

    # Calculate output shapes
    nb_rows_out = int(
        (nb_rows_in + 2 * padding[0] - dilation[0] * (kernel[0] - 1) - 1) / stride[0]
        + 1
    )
    nb_cols_out = int(
        (nb_cols_in + 2 * padding[1] - dilation[1] * (kernel[1] - 1) - 1) / stride[1]
        + 1
    )

    # Apply padding to the input
    if padding != (0, 0):
        a_sh = torch.nn.functional.pad(
            a_sh, (padding[1], padding[1], padding[0], padding[0]), "constant"
        )
        # Update shape after padding
        nb_rows_in += 2 * padding[0]
        nb_cols_in += 2 * padding[1]

    res = []
    # TODO: make this operation more efficient in order to be used with cnn modules.
    for batch in range(batch_size):
        for channel in range(nb_channels):
            for r_in in range(0, nb_rows_in - (kernel[0] - 1), stride[0]):
                for c_in in range(0, nb_cols_in - (kernel[1] - 1), stride[1]):
                    m, _ = maxpool(
                        a_sh[
                            batch,
                            channel,
                            r_in : r_in + kernel[0],
                            c_in : c_in + kernel[1],
                        ].child
                    )
                    res.append(m.wrap())

    res = torch.stack(res).reshape(batch_size, nb_channels, nb_rows_out, nb_cols_out)
    return res