Say your tensor x is of shape [batch, 1, 300, 300]
You can do it by dividing the task in a few steps : difference, then abs, then sum

Compute the differences, using a custom kernel that outputs the 8 differences, on 8 channels. This kernel would be of shape [8, 1, 3, 3] and would look something like this:


w = torch.FloatTensor([
[[-1, 0, 0], [0, 1, 0], [0, 0, 0]],
[[0, -1, 0], [0, 1, 0], [0, 0, 0]],
[[0, 0, -1], [0, 1, 0], [0, 0, 0]],
[[0, 0, 0], [-1, 1, 0], [0, 0, 0]],
[[0, 0, 0], [0, 1, -1], [0, 0, 0]],
[[0, 0, 0], [0, 1, 0], [-1, 0, 0]],
[[0, 0, 0], [0, 1, 0], [0, -1, 0]],
[[0, 0, 0], [0, 1, 0], [0, 0, -1]]
]).unsqueeze(1)

You can see that each of the 8 “filters” will compute the difference with one neighbor.
And with y = torch.nn.functional.conv2d(x, w) 7, the tensor y is of shape [batch, 8, 298, 298].


You want the absolute, so y = y.abs()


Do the sum over the channels dimension : y = y.sum(dim=1, keepdim=True)


You get the desired y with the shape [batch, 1, 298, 298]
I’m not sure if using Unfold / Fold instead would be relevant.