# Implementing a variation of MaxPool2d

Hi all, I need to implement a variation of MaxPool2d where all maximum values are selected and their indices are returned in case of equality instead of the first one and adapt MaxUnpool2d to it.

Here is an example,

``````pool = CustomMaxPool2d(2, stride=2, return_indices=True)
unpool = CustomMaxUnpool2d(2, stride=2)

inp = torch.tensor([[[[ 6.,  2,  3,  4],
[ 5,  6,  8,  8],
[ 9, 10, 11, 16],
[14, 14, 16, 16]]]])

output, indices = pool(inp)
print('maxPool:\n', output)
print('pool indices:\n', indices)
print('maxUnpool:\n', unpool(output, indices))

maxPool:
tensor([[[[ 6.,  8.],
[14., 16.]]]])

pool indices:
tensor([[[[[0, 5], [6, 7]],
[[12, 13], [11, 14, 15]]]]])

maxUnpool:
tensor([[[[ 6.,  0.,  0.,  0.],
[ 0.,  6.,  8.,  8.],
[ 0.,  0.,  0., 16.],
[14.,  14.,  16.,  16.]]]])
``````

Is there a way to do this without modifying the c++ code ?

I would recommend to start with a manual implementation in Python via `unfold` to get the patches and apply your custom pooling logic to each patch.
Note that the currently returned indices have a variable shape and would thus not be representable in a standard tensor but e.g. a `list` instead or a nested tensor (which is an experimental feature at the moment).
Once this approach is working properly, you could try to reimplement it as a custom layer in the C++/CUDA backend.

By unfold do you mean `torch.nn.unfold` or `tensor.unfold` ?

Both should work, but the `nn.Unfold` module would provide more built-in options such as padding dilation etc.