Ce document décrit la sémantique de broadcast de XLA.
Qu'est-ce que la diffusion ?
Le broadcasting consiste à faire en sorte que des tableaux de formes différentes aient 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 différents rangs 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 légèrement 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 le broadcast, XLA nécessite des spécifications de broadcasting explicites sur les opérations entre des tableaux de rangs différents. Elle est différente 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ées sur des tableaux sans spécification explicite de dimensions de broadcast. Une opération binaire par élément entre une valeur scalaire et un tableau consiste à appliquer l'opération avec la valeur scalaire à chaque élément du tableau. Par exemple, ajouter une valeur scalaire à une matrice revient à produire une matrice dans laquelle chaque élément est une somme de la valeur 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 en matière de diffusion 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 vous ne spécifiez pas de diffusion, cette opération n'est pas valide. Pour demander correctement l'ajout d'un vecteur de matrice, spécifiez que la dimension de broadcasting est (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 des lignes et que la dimension 1 représente des 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 : envisagez d'ajouter un vecteur à trois éléments (dimension (3)) à une matrice 3 x 3 (dimensions (3,3)). Pour cet exemple, la diffusion peut se faire 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 broadcasting 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 un cube cuboïde 2 x 3 x 4 et une matrice 3 x 4, un tuple de diffusion (1,2) signifie que la matrice doit correspondre aux dimensions 1 et 2 du cube.
Ce type de diffusion est utilisé dans les opérations binaires dans 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 doivent 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 avec P et N.
Diffuser des tableaux de rang similaire avec des dimensions dégénérées
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 le sont. 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 présents, la forme du résultat contient le maximum des deux entrées pour chaque indice de dimension.
Exemples :
- (2,1) et (2,3) sont diffusées sur (2,3).
- (1,2,5) et (7,2,5) sont diffusées sur (7,2,5).
- (7,2,5) et (7,1,5) sont diffusées sur (7,2,5).
- (7,2,5) et (7,2,6) sont incompatibles et ne peuvent pas être diffusés.
Un cas particulier est pris en charge : chacun des tableaux d'entrée présente 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 sur (2,3). Pour plus d'exemples, consultez la documentation NumPy sur la diffusion.
Composition de l'annonce 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 1x2 peuvent être additionnés en utilisant les dimensions de diffusion de la 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 dans laquelle la valeur M est choisie pour correspondre à la taille de dimension correspondante dans le tableau 1x2. Par conséquent, une matrice 4x2 est produite:
|1 1| + [5 6]
|2 2|
|3 3|
|4 4|
Ensuite, "dégénérer la diffusion de dimension" diffuse la dimension 0 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 utilisant des dimensions de diffusion de (1, 2). Tout d'abord, la matrice 1x2 est diffusée jusqu'au rang 3 en utilisant les dimensions de broadcast pour générer un tableau Mx1x2 intermédiaire, dans lequel la taille de dimension M est déterminée par la taille du plus grand opérande (le tableau 4x3x1), ce qui génère un tableau intermédiaire 4x1x2. La dimension M est la dimension 0 (la plus à gauche), car les dimensions 1 et 2 sont mappées avec les dimensions de la matrice 1 x 2 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.