Nested advanced indexing returns dimensions in the wrong order

[zdevito]

from functorch import vmap
import torch
e = torch.rand(7, 4)
idx = torch.LongTensor([0, 1, 2, 3, 4, 5]).view(2, 3, 1)

# simple reference implementation for comparison
def _fake_vmap(f, in_dims=0):
    def w(input):
        r = [f(input.select(in_dims, i)) for i in range(input.size(in_dims))]
        return torch.stack(r, in_dims)
    return w

def with_vmap(_vmap):
    def g(idx_):
        def f(e_):
            return e_[idx_]
        return _vmap(f, in_dims=1)(e)
    r = _vmap(_vmap(g))(idx)
    return r

a = with_vmap(vmap)
b = with_vmap(_fake_vmap)
print(a.size(), b.size())
assert a == b

Produces:

torch.Size([4, 2, 3, 1]) torch.Size([2, 3, 1, 4])
Traceback (most recent call last):
  File "v.py", line 25, in <module>
    assert a == b
RuntimeError: The size of tensor a (2) must match the size of tensor b (3) at non-singleton dimension 1