Get just the values of `input` which are at `indices`. Arguments: ctx: the autograd context object input: (b, ...) 2+ dimensional tensor indices: (num_idx) 1D tensor
(ctx, input: torch.Tensor, indices: torch.Tensor)
| 20 | class IndexFirstAxis(torch.autograd.Function): |
| 21 | @staticmethod |
| 22 | def forward(ctx, input: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: |
| 23 | """Get just the values of `input` which are at `indices`. |
| 24 | |
| 25 | Arguments: |
| 26 | ctx: the autograd context object |
| 27 | input: (b, ...) 2+ dimensional tensor |
| 28 | indices: (num_idx) 1D tensor |
| 29 | """ |
| 30 | ctx.save_for_backward(indices) |
| 31 | assert input.ndim >= 2 |
| 32 | ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] # type: ignore |
| 33 | second_dim = other_shape.numel() # product of sizes of all but first dimension |
| 34 | # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing. |
| 35 | return torch.gather( |
| 36 | rearrange(input, "b ... -> b (...)"), # (b, ...) -> (b, second_dim) |
| 37 | 0, |
| 38 | repeat(indices, "z -> z d", d=second_dim), # (indices,) -> (indices, second_dim) |
| 39 | ).reshape(-1, *other_shape) # (num_idx, ...) |
| 40 | |
| 41 | @staticmethod |
| 42 | def backward(ctx, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None]: |