🐛 Describe the bug

When calling torch.vmap over torch.index_select with a batched index, the result does not match Jax's results:

Expected

from jax import numpy as jnp
import jax
jax.vmap(lambda x, y: x[y], [None, 0])(jnp.arange(3), jnp.array([0, 2]))
# Array([0, 2], dtype=int32)

PyTorch

torch.vmap(lambda x, y: torch.index_select(x, 0, y), (None, 0))(torch.arange(3), torch.tensor([0, 2]))
# tensor([[0, 0, 0],
#              [2, 2, 2]])

Versions

PT '2.2.0.dev20231128'

cc @ezyang @gchanan @zou3519 @kadeng @Chillee @samdow @kshitij12345 @janeyx99

Yeah, this is broken. Here's a test case just using PyTorch:

import torch
x = torch.arange(3)
y = torch.tensor([0, 2])

def f(x, y):
    return torch.index_select(x, 0, y)

result = torch.vmap(f, (None, 0))(x, y)
# vmap should be equivalent to a for-loop + stack over the vmapped Tensor
expected = torch.stack([f(x, y[0]), f(x, y[1])])
print(result)
print(expected)

The first place to look is

pytorch/aten/src/ATen/functorch/BatchRulesScatterOps.cpp

Line 1250 in 78b9454

m.impl("index_select", index_select_decomp);

. If the index_select decomposition is correct, then it's possible the gather batching rule is wrong (

pytorch/aten/src/ATen/functorch/BatchRulesScatterOps.cpp

Line 1259 in 78b9454

VMAP_SUPPORT(gather, gather_batch_rule);

)