PyTorch torch.where : masques, dimensions et gradients

torch.where(condition, input, other) choisit les valeurs élément par élément : la fonction prend input lorsque la condition booléenne est vraie et other lorsqu'elle est fausse. Elle permet de remplacer des valeurs en conservant la structure d'un tenseur, de sélectionner les résultats de deux calculs ou d'appliquer des valeurs de remplacement à certaines lignes ou colonnes.

L'appel à un seul argument, torch.where(condition), renvoie des tenseurs d'indices. Ce guide teste les deux formes, montre comment les masques se diffusent entre dimensions et explique pourquoi un résultat fini ne garantit pas un gradient valide.

Les six exemples s'exécutent indépendamment sur CPU avec Python 3.12.14 et PyTorch 2.15.0+cpu. Les vérifications associées portent sur les valeurs, les dimensions, les erreurs attendues et les gradients. Elles ne mesurent pas les performances sur GPU.

Sélectionner des valeurs avec une condition booléenne

Commencez par une comparaison comme x > 0. Elle produit un tenseur booléen de mêmes dimensions que x. Ici, les éléments positifs conservent leur valeur ; tous les autres deviennent -9.0. Le deuxième appel produit des indicateurs entiers et la dernière expression combine deux comparaisons.

import torch

x = torch.tensor([[-3., 2., 0.], [4., -1., 5.]])
condition = x > 0
selected = torch.where(condition, x, -9.0)
flags = torch.where(condition, 1, 0)
between = (x > 0) & (x < 5)
print(selected.tolist())
print(flags.tolist())
print(between.tolist())

[[-9.0, 2.0, -9.0], [4.0, -9.0, 5.0]]
[[0, 1, 0], [1, 0, 1]]
[[False, True, False], [True, False, False]]

Zéro suit la branche fausse, car la comparaison exige une valeur strictement positive. L'entrée reste inchangée dans cet exemple. Pour combiner des conditions élément par élément, entourez chaque comparaison de parenthèses et utilisez & pour ET ou | pour OU. Les opérateurs Python and et or n'effectuent pas cette opération élément par élément.

La référence de torch.where (en anglais) décrit les deux formes d'appel. Avec trois arguments, la condition doit être booléenne et les valeurs peuvent provenir de tenseurs ou de scalaires. Un scalaire convient lorsque toutes les positions rejetées doivent recevoir la même valeur.

Diffuser un masque sur les lignes ou les colonnes

Pour une matrice de dimensions (2, 3), un masque de dimensions (3,) s'applique aux colonnes. Un masque de dimensions (2, 1) s'applique aux lignes. L'alignement dépend des dimensions de droite, pas du nom de la variable.

import torch

x = torch.tensor([[1., 2., 3.], [4., 5., 6.]])
row_mask = torch.tensor([True, False])[:, None]
fallback = torch.tensor([10., 20., 30.])
by_row = torch.where(row_mask, x, fallback)
column_mask = torch.tensor([True, False, True])
by_column = torch.where(column_mask, x, 0.0)
print(tuple(row_mask.shape), tuple(by_row.shape))
print(by_row.tolist())
print(by_column.tolist())

(2, 1) (2, 3)
[[1.0, 2.0, 3.0], [10.0, 20.0, 30.0]]
[[1.0, 0.0, 3.0], [4.0, 0.0, 6.0]]

La première ligne est conservée. La deuxième reçoit les trois valeurs de fallback. L'ajout de [:, None] transforme le masque de longueur deux en colonne de deux lignes et rend l'alignement voulu explicite. Dans le second appel, le masque de colonnes conserve les colonnes zéro et deux sur les deux lignes.

Les règles de broadcasting de PyTorch (en anglais) comparent les dimensions à partir de la droite. Des tailles égales ou une taille de un peuvent s'aligner ; les dimensions initiales absentes peuvent être ajoutées. Les trois opérandes doivent être compatibles, y compris other. Examinez leurs dimensions ensemble avant de les modifier.

Une matrice carrée peut masquer une erreur : un masque de longueur deux convient à un tenseur (2, 2), mais agit sur ses colonnes. Utilisez une matrice non carrée pour tester un masque de lignes. Dans une chaîne de préparation des données, vérifiez quel axe représente les échantillons avant de transmettre les tenseurs à un DataLoader PyTorch (guide en anglais).

Obtenir les indices avec torch.where(condition)

Pour cette entrée à deux dimensions, la forme à un argument renvoie deux tenseurs : les indices de ligne et les indices de colonne. Les éléments correspondants forment des paires de coordonnées. En utilisant ce tuple pour indexer l'entrée, vous récupérez les valeurs sélectionnées.

import torch

x = torch.tensor([[2, -1, 4], [0, 5, -2]])
condition = x > 0
rows, columns = torch.where(condition)
values = x[rows, columns]
filled = torch.where(condition, x, 0)
empty_indices = torch.where(x > 10)
print(rows.tolist(), columns.tolist())
print(values.tolist(), tuple(values.shape))
print(filled.tolist(), tuple(filled.shape))
print([part.tolist() for part in empty_indices])

[0, 0, 1] [0, 2, 1]
[2, 4, 5] (3,)
[[2, 0, 4], [0, 5, 0]] (2, 3)
[[], []]

Les coordonnées sont (0, 0), (0, 2) et (1, 1). Les valeurs extraites forment un tenseur unidimensionnel de longueur trois. L'appel à trois arguments conserve une matrice de deux lignes et trois colonnes, avec des zéros aux positions rejetées.

Si aucune valeur ne correspond, cette matrice produit toujours deux tenseurs d'indices, tous deux vides. Choisissez la forme selon ce que l'étape suivante attend : des coordonnées, les seules valeurs retenues ou un tenseur contenant des remplacements. La forme à un argument équivaut à torch.nonzero(condition, as_tuple=True) ; elle ne renvoie pas directement les données sélectionnées.

Diagnostiquer les erreurs de type et de dimensions

Les exemples suivants échouent volontairement. Ils affichent les classes d'exception plutôt que les messages complets, dont la formulation peut varier selon la version.

import torch

x = torch.ones(2, 3)
cases = {
    "integer condition": lambda: torch.where(
        torch.tensor([1, 0, 1]), x, 0.0),
    "row mask without axis": lambda: torch.where(
        torch.tensor([True, False]), x, 0.0),
    "incompatible other": lambda: torch.where(
        x > 0, x, torch.zeros(2, 4)),
}
errors = {}
for name, operation in cases.items():
    try:
        operation()
    except RuntimeError as exc:
        errors[name] = type(exc).__name__
        print(name + ": " + errors[name])
    else:
        raise AssertionError("Expected failure: " + name)

integer condition: RuntimeError
row mask without axis: RuntimeError
incompatible other: RuntimeError

Le premier cas fournit des indicateurs entiers alors que cet appel exige un masque booléen. Construisez une comparaison exprimant votre règle, comme flags != 0. Le deuxième tente d'aligner un vecteur de longueur deux sur une dernière dimension de taille trois ; ajoutez l'axe nécessaire pour cibler les lignes, comme précédemment. Le dernier donne quatre colonnes aux valeurs de remplacement alors que l'entrée en possède trois.

Ne modifiez pas les dimensions uniquement pour faire disparaître l'exception. Notez ce que chaque axe représente, déterminez si la condition concerne les échantillons ou les caractéristiques, puis vérifiez les valeurs obtenues sur une petite matrice. Des dimensions compatibles peuvent produire une sélection incorrecte.

Les gradients suivent les valeurs finies sélectionnées

Ici, le masque booléen est fixé. Chaque valeur sélectionnée contribue à la perte avec le poids correspondant. Le scalaire de remplacement est utilisé à trois positions ; son gradient additionne donc leurs poids.

import torch

condition = torch.tensor([[True, False, True],
                          [False, False, True]])
x = torch.tensor([[1., 2., 3.], [4., 5., 6.]],
                 dtype=torch.float64, requires_grad=True)
fallback = torch.tensor(10., dtype=torch.float64, requires_grad=True)
weights = torch.tensor([[1., 2., 3.], [4., 5., 6.]],
                       dtype=torch.float64)
selected = torch.where(condition, x, fallback)
loss = (selected * weights).sum()
loss.backward()
print(selected.tolist())
print(loss.item())
print(x.grad.tolist())
print(fallback.grad.item())

[[1.0, 10.0, 3.0], [10.0, 10.0, 6.0]]
156.0
[[1.0, 0.0, 3.0], [0.0, 0.0, 6.0]]
11.0

Le gradient de l'entrée vaut zéro aux positions rejetées et prend la valeur du poids aux positions retenues. Celui du scalaire de remplacement vaut 2 + 4 + 5 = 11. Les tests comparent ces résultats à un calcul scalaire indépendant et à des différences finies centrées.

Cet exemple n'apprend pas la décision booléenne. Une comparaison fournit un masque discret ; n'attendez pas d'elle qu'elle transmette un gradient utile à un seuil à apprendre. Pour l'inférence et le contrôle de l'enregistrement des gradients, consultez le guide sur eval, no_grad et inference_mode (en anglais).

Éviter le calcul invalide avant de sélectionner le résultat

Python évalue les arguments avant d'appeler torch.where. Dans torch.where(valid, x / denominator, 0.0), la division a déjà eu lieu. Sélectionner zéro ensuite peut cacher une opération invalide dans la sortie tout en laissant un gradient NaN.

import torch

denominator = torch.tensor([0., 2.])
valid = denominator != 0
unsafe_x = torch.tensor([3., 6.], requires_grad=True)
unsafe = torch.where(valid, unsafe_x / denominator, 0.0)
unsafe.sum().backward()
print(unsafe.tolist(), torch.isfinite(unsafe_x.grad).tolist())

safe_x = torch.tensor([3., 6.], requires_grad=True)
safe_denominator = torch.where(valid, denominator, 1.0)
safe = torch.where(valid, safe_x / safe_denominator, 0.0)
safe.sum().backward()
print(safe.tolist(), safe_x.grad.tolist())

[0.0, 3.0] [False, True]
[0.0, 3.0] [0.0, 0.5]

Les deux résultats de la passe avant sont [0.0, 3.0]. Le premier contrôle des gradients montre une valeur non finie. Le second calcul remplace les dénominateurs nuls par un avant la division, puis sélectionne explicitement zéro à ces positions. Son gradient est [0.0, 0.5].

Cette correction traite le cas présenté, avec des entrées finies et un dénominateur nul. Elle ne valide pas tous les NaN, infinis, logarithmes ou racines carrées. Chaque opération enregistrée par autograd doit recevoir des opérandes valides. Les notes sur autograd (en anglais) expliquent pourquoi masquer un résultat invalide après une division peut encore dégrader les gradients.

Quelle opération de sélection choisir ?

• Utilisez torch.where(mask, x, fallback) pour choisir entre des valeurs dans une sortie obtenue par broadcasting.
• Utilisez torch.where(mask) pour obtenir des coordonnées.
• Utilisez l'indexation booléenne, comme x[mask], pour extraire les éléments correspondants de la matrice présentée ici avec un masque de mêmes dimensions.
• Utilisez une opération à indices entiers, comme torch.gather, lorsque les positions sont déjà définies numériquement le long d'un axe.

Avant d'intégrer une sélection à une boucle d'entraînement, vérifiez les dimensions de sortie, quelques valeurs exactes et les gradients des branches qui effectuent une division ou une autre opération à domaine restreint. Ces contrôles détectent des erreurs qu'un résultat apparemment plausible peut cacher.