torch.gather behavior changed from 1.5.1 to master

[zou3519]

🐛 Bug

The documentation states that for torch.gather(input, dim, index), the index tensor must have the same size in all dimensions as input, except for dimension dim. This check is respected in PyTorch 1.5.1, but not on master.

To Reproduce

On 1.5.1:

>>> t = torch.tensor([[1,2],[3,4]])
>>> index = torch.tensor([[0]])
>>> torch.gather(t, 1, index)
RuntimeError: Size does not match at dimension 0 get 2 vs 1

On master:

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

Expected behavior

I'm not sure. Either we should update the documentation or add an error check.

[gchanan]

bisect points here: #34690

[gchanan]

that change seems okay, because the relationship is still well defined, i.e. for dim==1 we have
out[i][j] = input[i][index[i][j]]
so given that index.shape == out.shape and that the shape out the output is smaller in the non-dim dimensions that's a well defined relationship.

[gchanan]

this seems strange, though:

torch.gather(torch.randn(3,3), 0, torch.tensor([0]))
tensor([-0.1960])

so you'd have...
out[i] = input[index[i]] but input[index[i]] is not a value.

[gchanan]

should probably do similar checks for scatter.

[gchanan]

other strangeness:

>>> x
tensor([[-1.1485, -0.3051, -0.2253],
        [ 0.4836, -0.7794, -1.8069],
        [-0.5331,  1.0695, -2.0397]])
>>> torch.gather(x, 0, torch.tensor([[[0]]]))
tensor([[[-1.1485]]])
>>> torch.gather(x, 1, torch.tensor([[[0]]]))
tensor([[[-1.1485]]])
>>> torch.gather(x, 2, torch.tensor([[[0]]]))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
IndexError: Dimension out of range (expected to be in range of [-2, 1], but got 2)

if we have:

out(0,0,0) = input(index(0,0,0), 0, 0) OR
out(0,0,0) = input(0, index(0,0,0), 0) OR
out(0,0,0) = input(0, 0, index(0,0,0)))

why is dim=2 not allowed? it doesn't follow indexing semantics anyway.

[gchanan]

from bisecting, it looks like this starts being accepted in #36875. Probably because it's an optimization that considered skipping over the TORCH_INTERNAL_ASSERT to be okay.

[ngimel]

#40359 is also related. @nikitaved can you please tighten the shape checks so that requirements for index dimensions are enforced?

[nikitaved]

I have included the checks for equal dimensions. The gather behavior on inputs with index.size(i) <= input.size(i), i != dim is well-defined and is similar to what scatter is doing. Shall we update the docs?

[gchanan]

@nikitaved where you have included the checks for equal dimensions?

[nikitaved]

Ah, sorry, here it goes: #41672.

[gchanan]

I think I did the previous gather examples from the commits I linked instead of the release branch.

>>> x
tensor([[-0.9640, -1.2366, -1.1208],
        [ 1.8341,  0.6243, -0.5042],
        [ 0.0528,  0.2580,  0.2960]])
# we don't check the dimensionality matches, but we do check it's not larger
>>> torch.gather(x, 0, torch.tensor(0))
tensor(-0.9640)
>>>torch.gather(x, 0, torch.tensor([0]))
tensor([-0.9640])
>>> torch.gather(x, 0, torch.tensor([[0]]))
tensor([[-0.9640]])
>>> torch.gather(x, 0, torch.tensor([[[0]]]))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
IndexError: Dimension out of range (expected to be in range of [-2, 1], but got 2)