Dilated layer takes more than `k` neightbours

The Dilated layer doesn’t take into account k. This can lead to taking more neighbours than intended.

t = torch.tensor([
    [0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1],
    [0, 1, 2, 3, 4, 5, 0, 1, 2, 3, 4, 5]
])

res = Dilated(k=2, dilation=2)(t)
print(res) # here 3 neighbours are taken even though the constructor specified 2.
# tensor([[0, 0, 0, 1, 1, 1],
#         [0, 2, 4, 0, 2, 4]])

[lightaime]

Thanks for the suggestion @zademn. That is definitely a good idea if we are dealing with a more complex case. But in our example, we always build knn graphs with k*d neighbors. To keep it simple, we prefer to leave it as it is.

[zademn]

A possible solution would be (using einops):

from einops import rearrange

t = torch.tensor(
    [
        [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2],
        [0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4],
    ]
)
k = 2
d = 2

u, counts = torch.unique(t[0], return_counts=True)
k_constructed = counts[0]  # assume we always find k neighbours. We can give this as a parameter too
res1 = rearrange(t, "e (n2 k_constructed) -> e n2 k_constructed", k_constructed=k_constructed)

# tensor([[[0, 0, 0, 0, 0],
#          [1, 1, 1, 1, 1],
#          [2, 2, 2, 2, 2]],

#         [[0, 1, 2, 3, 4],
#          [0, 1, 2, 3, 4],
#          [0, 1, 2, 3, 4]]])
res2 = res1[:, :, ::d]  # Res dilated
print(res2)
# tensor([[[0, 0, 0],
#          [1, 1, 1],
#          [2, 2, 2]],

#         [[0, 2, 4],
#          [0, 2, 4],
#          [0, 2, 4]]])
res3 = res2[:, :, :k] # Take first k neighbours
print(res3)
# tensor([[[0, 0],
#          [1, 1],
#          [2, 2]],

#         [[0, 2],
#          [0, 2],
#          [0, 2]]])
res4 = rearrange(res3, "e d1 d2 -> e (d1 d2)")
print(res4)
# tensor([[0, 0, 1, 1, 2, 2],
#         [0, 2, 0, 2, 0, 2]])