Ce document décrit la sémantique de broadcast de XLA.
Qu'est-ce que la diffusion ?
Le broadcasting est le processus qui consiste à créer des tableaux de formes différentes qui ont des formes compatibles pour les opérations arithmétiques. La terminologie est empruntée à NumPy broadcasting.
Le broadcasting 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'exemple de l'addition X+v, où X est une matrice (un tableau de rang 2) et v est un vecteur (un tableau de rang 1). Pour effectuer une addition par élément, XLA doit "diffuser" le vecteur v au même rang que la matrice X, en répliquant v un certain nombre de fois. La longueur du vecteur doit correspondre à au moins 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 la dimension du vecteur est (3). Le vecteur est diffusé en le répliquant sur des lignes pour obtenir:
|1 2 3| + |7 8 9| = |8 10 12|
|4 5 6| |7 8 9| |11 13 15|
Dans NumPy, c'est ce qu'on appelle la diffusion.
Principes
Le langage XLA est aussi strict et explicite que possible, évitant ainsi les caractéristiques "magiques" implicites. Ces caractéristiques peuvent faciliter la définition de certains calculs, mais au prix d'un plus grand nombre d'hypothèses intégrées dans le 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 la diffusion, XLA nécessite des spécifications de diffusion explicites sur les opérations entre des tableaux de rangs différents. Ceci est différent de NumPy, qui déduit la spécification lorsque cela est possible.
Diffuser un tableau de rang inférieur sur un tableau de rang supérieur
Les scalaires peuvent toujours être diffusés sur des tableaux sans spécification explicite des dimensions de broadcast. Une opération binaire par élément entre une valeur scalaire et un tableau implique d'appliquer l'opération avec la valeur scalaire à chaque élément du tableau. Par exemple, ajouter un scalaire à une matrice revient à produire une matrice dans laquelle chaque élément est une somme de l'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 broadcast peuvent être capturés en utilisant 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 higher-rank à faire correspondre au tableau lower-rank.
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). Si la diffusion n'est pas spécifiée, cette opération n'est pas valide. Pour demander correctement l'ajout d'un vecteur matriciel, spécifiez la dimension de diffusion sur (1), ce qui signifie que la dimension du vecteur est mise en correspondance avec la dimension 1 de la matrice. En 2D, si la dimension 0 représente les lignes et que la dimension 1 représente les colonnes, cela signifie que chaque élément du vecteur devient une colonne de taille correspondant au nombre de lignes de la matrice:
|7 8 9| ==> |7 8 9|
|7 8 9|
Pour prendre un exemple plus complexe, vous pouvez ajouter un vecteur à trois éléments (dimension (3)) à une matrice 3 x 3 (dimensions (3,3)). Pour cet exemple, la diffusion peut se produire de deux manières:
(1) Vous pouvez utiliser une dimension de diffusion de 1. Chaque élément vectoriel 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) Vous pouvez utiliser une dimension de diffusion de 0. Chaque élément vectoriel 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 broadcast peuvent être un tuple qui décrit comment une forme de rang plus petite est diffusée dans une forme de rang plus grande. Par exemple, pour une matrice cuboïde 2x3x4 et une matrice 3x4, un tuple de diffusion (1,2) signifie une mise en correspondance de la matrice avec les dimensions 1 et 2 du carré.
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 broadcasting permet de faire correspondre un tableau de rang inférieur à un tableau de rang supérieur en spécifiant les dimensions du tableau de rang supérieur à mettre en correspondance. 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 de rang supérieur. Les valeurs du vecteur sont ensuite diffusées depuis la dimension correspondante vers toutes les autres dimensions.
Pour faire correspondre une matrice TxV au tableau MxNxPxQ, une paire de dimensions de broadcast 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 être l'ordre dans lequel les dimensions du tableau de rang inférieur sont censées correspondre aux dimensions du tableau de rang supérieur. Le premier élément du tuple spécifie la dimension du tableau de rang supérieur qui doit correspondre à la dimension 0 du tableau de rang inférieur. Le deuxième élément du tuple spécifie la dimension du tableau de rang supérieur qui doit correspondre à la dimension 1 du tableau de rang inférieur, 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 de faire correspondre V à N et T à P. Il est également illégal de faire correspondre V à P et N.
Diffuser des tableaux de rang similaire avec des dimensions dégénérer
Un problème connexe consiste à diffuser deux tableaux ayant le même rang, 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 sont compatibles. Deux dimensions sont compatibles si:
- Ils sont égaux, ou
- L'une d'entre elles est 1 (dimension "dégénérée").
Lorsque deux tableaux compatibles sont rencontrés, la forme du résultat contient le maximum des deux entrées à chaque indice de dimension.
Exemples :
- (2,1) et (2,3) sont diffusés à (2,3).
- (1,2,5) et (7,2,5) sont diffusés à (7,2,5).
- (7,2,5) et (7,1,5) sont diffusés à (7,2,5).
- (7,2,5) et (7,2,6) sont incompatibles et ne peuvent pas être diffusés.
Il existe un cas particulier (qui est également pris en charge) dans lequel chacun des tableaux d'entrée a une dimension dégénérée à un indice différent. Dans ce cas, le résultat est une "opération externe": (2,1) et (1,3) sont diffusés à (2,3). Pour plus d'exemples, consultez la documentation de NumPy sur la diffusion.
Composition de diffusion
La diffusion d'un tableau de rang inférieur vers un tableau de rang supérieur et la diffusion à l'aide de dimensions dégénérées peuvent être effectuées dans la même opération binaire. Par exemple, un vecteur de taille 4 et une matrice de taille 1 x 2 peuvent être additionnés en utilisant les dimensions de broadcast de valeur (0):
|1 2 3 4| + [5 6] // [5 6] is a 1x2 matrix, not a vector.
Tout d'abord, le vecteur est diffusé jusqu'au rang 2 (matrice) en utilisant les dimensions de broadcast. 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 produit une matrice de taille 4xM, où la valeur M est choisie pour correspondre à la taille de dimension correspondante dans le tableau 1x2. Vous obtenez donc une matrice 4x2:
|1 1| + [5 6]
|2 2|
|3 3|
|4 4|
Ensuite, "dégénérer la diffusion de dimension" diffuse la dimension zéro de la matrice 1x2 pour correspondre à 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 compliqué est une matrice de taille 1x2 ajoutée à un tableau de taille 4x3x1 en utilisant des dimensions de diffusion de (1, 2). Tout d'abord, la matrice 1x2 est diffusée jusqu'au rang 3 à l'aide des dimensions de broadcast afin de produire un tableau Mx1x2 intermédiaire dans lequel la taille de 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. La dimension M est au niveau de 0 (la dimension la plus à gauche), car les dimensions 1 et 2 sont mappées avec les dimensions de la matrice 1x2 d'origine, comme les dimensions de diffusion sont (1, 2). Ce tableau intermédiaire peut être ajouté à la matrice 4x3x1 à l'aide de la diffusion de dimensions dégénérées pour produire un résultat de tableau 4x3x2.