Ce document décrit la sémantique de diffusion de XLA.
Qu'est-ce que la diffusion ?
Le broadcasting est le processus qui permet aux tableaux de formes différentes d'avoir des formes compatibles pour les opérations arithmétiques. La terminologie est empruntée à la diffusion NumPy.
La diffusion peut être nécessaire pour les opérations entre des tableaux multidimensionnels de rangs différents ou entre des tableaux multidimensionnels de formes différentes, mais compatibles. Prenons l'addition X+v où X est une matrice (un tableau à deux dimensions) et v est un vecteur (un tableau à une dimension). Pour effectuer une addition élément par élément, XLA doit "diffuser" le vecteur v au même nombre de dimensions que la matrice X, en répliquant v un certain nombre de fois. La longueur du vecteur doit correspondre à au moins l'une des dimensions de la matrice.
Exemple :
|1 2 3| + |7 8 9|
|4 5 6|
Les dimensions de la matrice sont (2,3) et celles du vecteur sont (3). Le vecteur est diffusé en le répliquant sur les lignes pour obtenir :
|1 2 3| + |7 8 9| = |8 10 12|
|4 5 6| |7 8 9| |11 13 15|
Dans NumPy, cette opération est appelée diffusion.
Principes
Le langage XLA est aussi strict et explicite que possible, en évitant les fonctionnalités "magiques" implicites. Ces fonctionnalités peuvent faciliter légèrement la définition de certains calculs, mais au prix de davantage d'hypothèses intégrées au code utilisateur, qui seront difficiles à modifier à long terme. Si nécessaire, des fonctionnalités magiques implicites peuvent être ajoutées dans des wrappers au niveau du client.
En ce qui concerne le broadcasting, XLA exige des spécifications de broadcasting explicites sur les opérations entre des tableaux de différents rangs. Cela diffère de NumPy, qui infère la spécification lorsque cela est possible.
Diffuser un tableau de dimension inférieure sur un tableau de dimension supérieure
Les scalaires peuvent toujours être diffusés sur des tableaux sans spécification explicite des dimensions de diffusion. Une opération binaire par élément entre un scalaire et un tableau signifie que l'opération avec le scalaire est appliquée à chaque élément du tableau. Par exemple, ajouter un scalaire à une matrice signifie produire une matrice dans laquelle chaque élément est la somme du scalaire et de l'élément correspondant de la matrice d'entrée.
|1 2 3| + 7 = |8 9 10|
|4 5 6| |11 12 13|
La plupart des besoins de diffusion peuvent être capturés à l'aide d'un tuple de dimensions sur une opération binaire. Lorsque les entrées de l'opération ont des rangs différents, ce tuple de diffusion spécifie la ou les dimensions du tableau de dimension supérieure à faire correspondre avec le tableau de dimension inférieure.
Prenons l'exemple précédent. Au lieu d'ajouter un scalaire à une matrice (2,3), ajoutez un vecteur de dimension (3) à une matrice de dimensions (2,3). Cette opération n'est pas valide si vous ne spécifiez pas la diffusion. Pour demander correctement l'addition matrice-vecteur, spécifiez que la dimension de diffusion est (1), ce qui signifie que la dimension du vecteur correspond à la dimension 1 de la matrice. En 2D, si la dimension 0 représente les lignes et la dimension 1 les colonnes, cela signifie que chaque élément du vecteur devient une colonne dont la taille correspond au nombre de lignes de la matrice :
|7 8 9| ==> |7 8 9|
|7 8 9|
Prenons un exemple plus complexe : ajoutons un vecteur à trois éléments (dimension (3)) à une matrice 3x3 (dimensions (3,3)). Dans cet exemple, la diffusion peut se faire de deux manières :
(1) Une dimension de diffusion de 1 peut être utilisée. Chaque élément de vecteur devient une colonne et le vecteur est dupliqué pour chaque ligne de la matrice.
|7 8 9| ==> |7 8 9|
|7 8 9|
|7 8 9|
(2) Une dimension de diffusion de 0 peut être utilisée. Chaque élément de vecteur devient une ligne et le vecteur est dupliqué pour chaque colonne de la matrice.
|7| ==> |7 7 7|
|8| |8 8 8|
|9| |9 9 9|
Les dimensions de diffusion peuvent être un tuple qui décrit comment une forme de dimension inférieure est diffusée dans une forme de dimension supérieure. Par exemple, étant donné un cuboïde 2x3x4 et une matrice 3x4, un tuple de diffusion (1,2) signifie que la matrice correspond aux dimensions 1 et 2 du cuboïde.
Ce type de diffusion est utilisé dans les opérations binaires de XlaBuilder, si l'argument broadcast_dimensions est fourni. Par exemple, consultez XlaBuilder::Add.
Dans le code source XLA, ce type de diffusion est parfois appelé diffusion "InDim".
Définition formelle
L'attribut de diffusion permet de faire correspondre un tableau de dimension inférieure à un tableau de dimension supérieure en spécifiant les dimensions du tableau de dimension supérieure à faire correspondre. Par exemple, pour un tableau de dimensions MxNxPxQ, un vecteur de dimension T peut être mis en correspondance comme suit :
MxNxPxQ
dim 3: T
dim 2: T
dim 1: T
dim 0: T
Dans chaque cas, T doit être égal à la dimension correspondante du tableau multidimensionnel. Les valeurs du vecteur sont ensuite diffusées de la dimension correspondante à toutes les autres dimensions.
Pour faire correspondre une matrice TxV au tableau MxNxPxQ, une paire de dimensions de diffusion est utilisée :
MxNxPxQ
dim 2,3: T V
dim 1,2: T V
dim 0,3: T V
etc...
L'ordre des dimensions dans le tuple de diffusion doit correspondre à l'ordre dans lequel les dimensions du tableau de dimension inférieure doivent correspondre à celles du tableau de dimension supérieure. Le premier élément du tuple spécifie la dimension du tableau de dimension supérieure qui doit correspondre à la dimension 0 du tableau de dimension inférieure. Le deuxième élément du tuple spécifie la dimension du tableau de dimension supérieure qui doit correspondre à la dimension 1 du tableau de dimension inférieure, et ainsi de suite. L'ordre des dimensions de diffusion doit être strictement croissant. Par exemple, dans l'exemple précédent, il est illégal d'associer V à N et T à P. Il est également illégal d'associer V à P et N.
Diffuser des tableaux de dimensions similaires avec des dimensions dégénérées
Un problème connexe consiste à diffuser deux tableaux qui ont le même nombre de dimensions, mais des tailles de dimensions différentes. Comme avec NumPy, cela n'est possible que lorsque les tableaux sont compatibles. Deux tableaux sont compatibles lorsque toutes leurs dimensions le sont. Deux dimensions sont compatibles si :
- Elles sont égales.
- L'une d'elles est 1 (dimension "dégénérée").
Lorsque deux tableaux compatibles sont rencontrés, la forme du résultat correspond au maximum des deux entrées à chaque index de dimension.
Exemples :
- (2,1) et (2,3) diffusent vers (2,3).
- (1,2,5) et (7,2,5) sont diffusés sur (7,2,5).
- (7,2,5) et (7,1,5) sont diffusés sur (7,2,5).
- (7,2,5) et (7,2,6) sont incompatibles et ne peuvent pas être diffusés.
Un cas particulier se produit (et est également pris en charge) lorsque chacun des tableaux d'entrée comporte une dimension dégénérée à un index différent. Dans ce cas, le résultat est une "opération externe" : (2,1) et (1,3) sont diffusés sur (2,3). Pour obtenir d'autres exemples, consultez la documentation NumPy sur le broadcasting.
Composition de la diffusion
La diffusion d'un tableau de dimension inférieure vers un tableau de dimension supérieure et la diffusion à l'aide de dimensions dégénérées peuvent toutes deux être effectuées dans la même opération binaire. Par exemple, un vecteur de taille 4 et une matrice de taille 1x2 peuvent être ajoutés à l'aide de dimensions de diffusion de valeur (0) :
|1 2 3 4| + [5 6] // [5 6] is a 1x2 matrix, not a vector.
Le vecteur est d'abord diffusé jusqu'à deux dimensions (matrice) à l'aide des dimensions de diffusion. La valeur unique (0) dans les dimensions de diffusion indique que la dimension zéro du vecteur correspond à la dimension zéro de la matrice. Cela génère une matrice de taille 4xM, où la valeur M est choisie pour correspondre à la taille de dimension correspondante dans le tableau 1x2. Une matrice de 4 x 2 est donc produite :
|1 1| + [5 6]
|2 2|
|3 3|
|4 4|
La "diffusion de dimension dégénérée" diffuse ensuite la dimension zéro de la matrice 1x2 pour qu'elle corresponde à la taille de dimension correspondante du côté droit :
|1 1| + |5 6| |6 7|
|2 2| + |5 6| = |7 8|
|3 3| + |5 6| |8 9|
|4 4| + |5 6| |9 10|
Un exemple plus complexe est une matrice de taille 1x2 ajoutée à un tableau de taille 4x3x1 à l'aide de dimensions de diffusion (1, 2). La matrice 1x2 est d'abord diffusée jusqu'à trois dimensions à l'aide des dimensions de diffusion pour produire un tableau intermédiaire Mx1x2, où la taille de la dimension M est déterminée par la taille de l'opérande le plus grand (le tableau 4x3x1), ce qui produit un tableau intermédiaire 4x1x2. M se trouve à la dimension 0 (la dimension la plus à gauche), car les dimensions 1 et 2 sont mappées sur les dimensions de la matrice 1x2 d'origine, car les dimensions de diffusion sont (1, 2). Ce tableau intermédiaire peut être ajouté à la matrice 4x3x1 à l'aide de la diffusion des dimensions dégénérées pour produire un tableau 4x3x2.