PyTorch torch.gather : indices, dimensions et gradients

torch.gather(input, dim, index) sélectionne des valeurs d’un tenseur le long d’un axe à partir de positions entières. Le tenseur d’indices peut choisir des positions différentes pour chaque ligne ou élément du lot, et la sortie prend la forme de ce tenseur d’indices. L’opération récupère des valeurs existantes ; elle ne les additionne pas et ne les trie pas automatiquement.

Ce guide commence par une matrice à deux lignes, puis sélectionne des vecteurs complets dans un lot de séquences. Il compare aussi des opérations d’indexation proches, reproduit des erreurs courantes et vérifie les gradients lorsqu’une valeur source est sélectionnée plusieurs fois.

Les six exemples indépendants s’exécutent avec Python 3.12.14 et PyTorch 2.15.0+cpu. Nous les avons validés par 47 contrôles, dont des calculs sur des listes Python et une comparaison de gradients par différences finies. La référence officielle gather décrit l’API ; les résultats présentés proviennent de nos exemples sur CPU.

Choisir l’axe dont les coordonnées changent

Pour une matrice avec dim=1, chaque indice choisit une colonne dans la même ligne. Avec dim=0, il choisit une ligne dans la même colonne. Suivez une coordonnée de sortie avant d’interpréter le tenseur entier :

import torch

x = torch.tensor([[10, 11, 12], [20, 21, 22]])
columns = torch.tensor([[2, 0], [1, 1]], dtype=torch.long)
rows = torch.tensor([[1, 0, 1], [0, 1, 0]], dtype=torch.long)
print('Select columns:', torch.gather(x, 1, columns).tolist())
print('Select rows:', torch.gather(x, 0, rows).tolist())

Select columns: [[12, 10], [21, 21]]
Select rows: [[20, 11, 22], [10, 21, 12]]

Dans l’exemple des colonnes, la première ligne demande les colonnes 2 et 0, donnant [12, 10]. La seconde demande deux fois la colonne 1, donnant [21, 21]. Les indices désignent des positions, pas les nombres à rechercher dans l’entrée.

Pour la sélection des lignes, la position de sortie (0, 0) utilise l’indice 1 et lit x[1, 0], soit 20. La position (0, 1) utilise l’indice 0 et lit x[0, 1], soit 11. Seule la coordonnée sur dim est remplacée ; l’autre reste attachée à la position de sortie.

Changer dim en conservant le même tenseur d’indices peut donc modifier le sens de l’appel et sa validité. Nos tests comparent les deux sorties à des accès indépendants dans des listes Python. Ils confirment aussi que dim=-1 choisit le dernier axe de cette matrice, comme dim=1.

Déduire la forme de sortie de index

Un tenseur d’indices plus petit ne se répète pas automatiquement sur l’entrée. Ici, l’entrée possède deux lignes, mais les indices n’en ont qu’une :

import torch

x = torch.tensor([[10, 11, 12], [20, 21, 22]])
index = torch.tensor([[2, 0, 1, 2]], dtype=torch.long)
selected = torch.gather(x, 1, index)
print('One index row:', selected.tolist())
print('Output shape:', tuple(selected.shape))
shared = index.expand(x.shape[0], -1)
print('Explicitly expanded:', torch.gather(x, 1, shared).tolist())
empty = torch.empty((2, 0), dtype=torch.long)
print('Empty output shape:', tuple(torch.gather(x, 1, empty).shape))

One index row: [[12, 10, 11, 12]]
Output shape: (1, 4)
Explicitly expanded: [[12, 10, 11, 12], [22, 20, 21, 22]]
Empty output shape: (2, 0)

Le premier résultat contient uniquement les sélections de la première ligne d’entrée. Aucune seconde ligne n’est créée implicitement. L’appel explicite à expand fournit les quatre mêmes choix de colonnes aux deux lignes, ce qui donne le résultat à deux lignes affiché ensuite.

Les quatre colonnes demandées montrent aussi que la sortie peut être plus longue que la source sur l’axe sélectionné. La source n’a que trois colonnes, mais les positions valides peuvent se répéter. Chaque indice doit toujours désigner une colonne existante.

Pour des indices non vides, input et index doivent avoir le même nombre d’axes. Sur les axes autres que dim, la taille de index ne doit pas dépasser celle de input. Le premier résultat respecte ces règles : une ligne d’indices tient dans les deux lignes de l’entrée. Trois lignes d’indices provoqueraient une erreur.

Les indices vides produisent une sortie de forme (2, 0). Traitez ce cas séparément lors des tests d’un pipeline. La réussite d’un appel avec des indices vides ne prouve pas que la même construction satisfera les contrôles appliqués à des indices non vides.

Sélectionner des vecteurs complets dans chaque séquence

Supposons que features ait pour forme (lot, temps, caractéristique) et que chaque élément du lot demande ses propres positions temporelles. La liste des positions a pour forme (lot, temps_sélectionné). Pour récupérer toutes les caractéristiques à chaque position choisie, ajoutez l’axe des caractéristiques et étendez les positions sur cet axe :

import torch

features = torch.arange(12).reshape(2, 3, 2)  # batch, time, feature
positions = torch.tensor([[2, 0], [1, 1]], dtype=torch.long)
index = positions.unsqueeze(-1).expand(-1, -1, features.shape[-1])
selected = torch.gather(features, dim=1, index=index)
print('Index shape:', tuple(index.shape))
print('Selected vectors:', selected.tolist())

Index shape: (2, 2, 2)
Selected vectors: [[[4, 5], [0, 1]], [[8, 9], [8, 9]]]

La première séquence sélectionne la position 2 puis la position 0, donnant [[4, 5], [0, 1]]. La seconde sélectionne deux fois la position 1, donnant [[8, 9], [8, 9]]. Répéter le même indice temporel sur les deux coordonnées de caractéristiques permet de conserver chaque vecteur entier.

Les indices et la sortie ont tous deux pour forme (2, 2, 2). L’opération conserve l’ordre demandé par positions ; elle ne rétablit pas l’ordre chronologique et ne supprime pas les doublons. Notre vérification compare ces vecteurs à une sélection directe dans des listes Python.

Dans un pipeline d’entraînement, vérifiez que positions correspond bien au même lot que features. Les dimensions seules ne détectent pas des indices associés par erreur à un autre lot. Notre guide PyTorch DataLoader, en anglais, traite du chargement et de la constitution des lots qui fournissent ces entrées.

Utiliser index_select pour des positions communes

Lorsque toutes les lignes doivent prendre les mêmes colonnes, une liste d’indices à un axe avec index_select exprime directement ce choix. Gather peut donner le même résultat après extension explicite de cette liste commune :

import torch

x = torch.tensor([[10, 11, 12], [20, 21, 22]])
common = torch.tensor([2, 0], dtype=torch.long)
print('Shared columns:', torch.index_select(x, 1, common).tolist())
expanded = common.unsqueeze(0).expand(x.shape[0], -1)
print('Gather agrees:', torch.equal(
    torch.index_select(x, 1, common), torch.gather(x, 1, expanded)
))
scores = torch.tensor([[0.2, 0.8, 0.1], [0.9, 0.3, 0.4]], dtype=torch.float64)
best = scores.argmax(dim=1, keepdim=True)
print('Best indices:', best.tolist())
print('Best scores:', torch.gather(scores, 1, best).tolist())
print('take_along_dim agrees:', torch.equal(
    torch.gather(scores, 1, best), torch.take_along_dim(scores, best, dim=1)
))

Shared columns: [[12, 10], [22, 20]]
Gather agrees: True
Best indices: [[1], [0]]
Best scores: [[0.8], [0.9]]
take_along_dim agrees: True

Les deux appels avec colonnes communes renvoient les colonnes 2 et 0 pour chaque ligne. Le premier exemple gather choisissait plutôt [2, 0] pour une ligne et [1, 1] pour l’autre. Ce choix propre à chaque ligne justifie l’emploi d’une matrice d’indices.

La seconde partie sélectionne le score maximal de chaque ligne. Conserver l’axe dans argmax donne à best la forme (2, 1), compatible avec gather sans ajout ultérieur d’un axe. Les scores sélectionnés sont 0.8 et 0.9.

take_along_dim est également conçu pour utiliser les indices renvoyés par des opérations comme argmax et argsort. Il donne le même résultat que gather dans cet exemple précis. N’en déduisez pas que toutes les formes acceptées ou tous les arguments par défaut sont interchangeables.

Reproduire les erreurs d’indexation courantes

Utilisez des indices torch.long dans ces exemples. Examinez leur forme et leurs valeurs avant l’appel à gather. Ces cinq cas enfreignent volontairement des conditions différentes :

import torch

x = torch.tensor([[10, 11, 12], [20, 21, 22]])
cases = {
    'Floating-point indices': torch.tensor([[0., 1.], [1., 2.]]),
    'Missing index axis': torch.tensor([0, 1]),
    'Negative index value': torch.tensor([[0, -1], [1, 2]]),
    'Index beyond final column': torch.tensor([[0, 3], [1, 2]]),
    'Too many index rows': torch.zeros((3, 1), dtype=torch.long),
}
for name, index in cases.items():
    try:
        torch.gather(x, dim=1, index=index)
    except RuntimeError:
        print(name + ': RuntimeError')

Floating-point indices: RuntimeError
Missing index axis: RuntimeError
Negative index value: RuntimeError
Index beyond final column: RuntimeError
Too many index rows: RuntimeError

Le tenseur flottant n’a pas le bon dtype. Le tenseur à un axe manque de l’axe des lignes. La valeur négative et la valeur 3 sortent des positions de colonne valides, qui sont 0, 1 et 2. Le dernier cas comporte plus de lignes que l’entrée sur un axe qui n’est pas sélectionné.

Un dim négatif et une valeur d’indice négative n’ont pas le même sens. Le test dim=-1 réussit parce qu’il désigne le dernier axe. La valeur d’indice -1 ne donne pas accès à la dernière colonne comme dans certaines indexations Python ; notre appel gather déclenche une erreur.

Le code affiche RuntimeError sans reproduire les messages détaillés, dont la formulation peut varier selon la version. Sur GPU, vérifiez aussi les périphériques du tenseur source et des indices. Cet article teste le comportement CPU et ne prétend pas que CUDA signale les erreurs de la même manière ou au même moment.

Suivre les sélections répétées pendant backward

Gather peut participer à autograd lorsque la source contient des valeurs flottantes qui demandent des gradients. Répéter une position source répète sa contribution à la perte. Cette somme pondérée rend les contributions visibles :

import torch

x = torch.tensor([[1., 2., 3.], [4., 5., 6.]],
                 dtype=torch.float64, requires_grad=True)
index = torch.tensor([[2, 2, 0], [1, 1, 1]], dtype=torch.long)
weights = torch.tensor([[1., 2., 4.], [1., 3., 5.]], dtype=torch.float64)
selected = torch.gather(x, 1, index)
loss = (selected * weights).sum()
loss.backward()
print('Selected:', selected.tolist())
print('Weighted loss:', loss.item())
print('Input gradient:', x.grad.tolist())
print('Index requires_grad:', index.requires_grad)

Selected: [[3.0, 3.0, 1.0], [5.0, 5.0, 5.0]]
Weighted loss: 58.0
Input gradient: [[4.0, 0.0, 3.0], [0.0, 9.0, 0.0]]
Index requires_grad: False

Dans la première ligne, la colonne source 2 est sélectionnée deux fois avec les poids 1 et 2 : son gradient vaut 3. La colonne source 0 est sélectionnée une fois avec le poids 4. La colonne 1 n’est pas utilisée et reçoit zéro. Dans la seconde ligne, la colonne 1 reçoit les trois poids : 1 + 3 + 5 = 9.

La perte pondérée vaut 58. Nos tests comparent toute la matrice de gradients à ces comptes indépendants, puis à des différences finies centrées d’une perte scalaire Python. Ils vérifient ainsi l’accumulation, y compris les gradients nuls aux positions non sélectionnées.

Le tenseur d’indices entiers ne porte pas de gradients. Si un indice provient d’un choix discret comme argmax, cet exemple ne démontre pas un chemin différentiable à travers ce choix. Il vérifie les gradients par rapport aux valeurs sources sélectionnées.

L’enregistrement des gradients est une décision distincte de la sélection des valeurs. Notre guide de eval, no_grad et inference_mode, en anglais, présente les contrôles utilisés pour l’entraînement et l’inférence. Ces tests portent sur des gradients denses sur CPU ; sparse_grad, les performances CUDA et la précision mixte demandent leur propre validation.

Vérifier les coordonnées avant de changer d’échelle

Notez les axes de la source, l’axe sélectionné et la forme de sortie voulue. Suivez ensuite une valeur de sortie jusqu’à sa coordonnée source. Pour des séquences en lot, vérifiez que chaque position demandée s’étend sur l’axe de caractéristiques à conserver. Un petit test de valeurs exactes peut révéler une erreur d’indexation qu’une assertion sur les dimensions laisserait passer.