Is`_index_first_axis` functionally equal to FA2's `index_first_axis`

[A1waysBeenHere]

Ref to the #40002 . This PR consoliated this function into transformers.modeling_falsh_attention_utils and removed the index_first_axis function that directly copied from flash-attention repo. I am wondering if they are functionally equally.

# The shape is recovered to the original one at the end, and ref to the comment, 
# it actually equal to return input[indices]
class IndexFirstAxis(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, indices):
        ctx.save_for_backward(indices)
        assert input.ndim >= 2
        ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:]
        second_dim = other_shape.numel()
        # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing.
        # return input[indices]
        return torch.gather(
            rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim)
        ).reshape(-1, *other_shape)


# The currently one, I don't think the shape are exactly same here, cuz no shape recoverd
def _index_first_axis(tensor, indices):
    """
    A local implementation of the PyTorch indexing operation `tensor[indices]` on the first axis,
    after flattening the first two dimensions of the tensor. This is functionally equivalent to
    FA2's `index_first_axis` and replaces the need to import it.
    """
    # The input tensor is expected to be of shape (batch, seq_len, ...). We flatten the first
    # two dimensions to get (total_tokens, ...) before indexing.
    reshaped_tensor = tensor.reshape(-1, *tensor.shape[2:])
    return reshaped_tensor[indices]

[vasqu]

Hey @A1waysBeenHere,

• short answer: Yes, they are equal.
• long answer: The reshape is necessary due to torch gather consolidating different shapes (see the rearrange, repeat shapes). You can also see Tri Dao's comment in the original implementation (right before the gather): "# TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. # return input[indices]" (assuming the first two dimensions are flattened, i.e. bsz, seq_len -> total_nnz)

Edit: Depending on how you used it, there may be a difference in the original shape. The new pure torch version expects a non-flattened tensor, while the other original one expects the flattened one already. This depends on the context of the pad/unpad functions. So it is more likely correct to say that they are equal under unpad and expect different initial shapes (bzs, seq_len, ...) vs (total_nnz, ...)

[vasqu]

Since it was originally not meant to be public-facing, I'm open for a PR here so that the functions are really equal, i.e. the initial shapes are the same (should be only the reshape to be moved outside).