A continuación, se describe la semántica de las operaciones definidas en la interfaz de XlaBuilder
. Por lo general, estas operaciones se asignan uno a uno a las operaciones definidas en la interfaz de RPC en xla_data.proto
.
Nota sobre la nomenclatura: El tipo de datos generalizado con el que se encarga XLA es un array de N dimensiones que contiene elementos de algún tipo uniforme (como un número de punto flotante de 32 bits). En toda la documentación, array se usa para denotar un array de dimensiones arbitrarias. Para mayor comodidad, los casos especiales tienen nombres más específicos y conocidos. Por ejemplo, un vector es un array de 1 dimensión y una matriz es un array de 2 dimensiones.
AfterAll
Consulta también XlaBuilder::AfterAll
.
AfterAll toma una cantidad variable de tokens y produce un solo token. Los tokens son tipos primitivos que se pueden agrupar en subprocesos entre operaciones de efectos secundarios para aplicar el orden. AfterAll
se puede usar como una unión de tokens para pedir una operación después de una operación set.
AfterAll(operands)
Argumentos | Tipo | Semántica |
---|---|---|
operands |
XlaOp |
cantidad variable de tokens |
AllGather
Consulta también XlaBuilder::AllGather
.
Realiza la concatenación entre réplicas.
AllGather(operand, all_gather_dim, shard_count, replica_group_ids,
channel_id)
Argumentos | Tipo | Semántica |
---|---|---|
operand
|
XlaOp
|
Es un array para concatenar entre réplicas. |
all_gather_dim |
int64 |
Dimensión de concatenación |
replica_groups
|
vector de vectores de int64 |
Grupos entre los que se realiza la concatenación |
channel_id
|
int64 opcional
|
ID de canal opcional para la comunicación entre módulos |
replica_groups
es una lista de grupos de réplicas entre los que se realiza la concatenación (el ID de réplica de la réplica actual se puede recuperar conReplicaId
). El orden de las réplicas en cada grupo determina el orden en el que se ubican sus entradas en el resultado.replica_groups
debe estar vacío (en ese caso, todas las réplicas pertenecen a un solo grupo, ordenadas de0
aN - 1
) o contener la misma cantidad de elementos que la cantidad de réplicas. Por ejemplo,replica_groups = {0, 2}, {1, 3}
realiza la concatenación entre las réplicas0
y2
, y1
y3
.shard_count
es el tamaño de cada grupo de réplicas. Esto lo necesitamos en los casos en que los objetosreplica_groups
están vacíos.channel_id
se usa para la comunicación entre módulos: solo las operacionesall-gather
con el mismochannel_id
pueden comunicarse entre sí.
La forma de la salida es la de entrada con el all_gather_dim
hecho shard_count
veces más grande. Por ejemplo, si hay dos réplicas y el operando tiene el valor [1.0, 2.5]
y [3.0, 5.25]
respectivamente en ambas réplicas, el valor de salida de esta op en la que all_gather_dim
es 0
será [1.0, 2.5, 3.0,
5.25]
en ambas réplicas.
AllReduce
Consulta también XlaBuilder::AllReduce
.
Realiza un procesamiento personalizado en las réplicas.
AllReduce(operand, computation, replica_group_ids, channel_id)
Argumentos | Tipo | Semántica |
---|---|---|
operand
|
XlaOp
|
un array o una tupla no vacía de arrays para reducir en las réplicas |
computation |
XlaComputation |
Cálculo de reducción |
replica_groups
|
vector de vectores de int64 |
Grupos entre los que se realizan las reducciones |
channel_id
|
int64 opcional
|
ID de canal opcional para la comunicación entre módulos |
- Cuando
operand
es una tupla de arreglos, se realiza el All-reduce en cada elemento de la tupla. replica_groups
es una lista de grupos de réplicas entre los que se realiza la reducción (el ID de réplica de la réplica actual se puede recuperar conReplicaId
).replica_groups
debe estar vacío (en cuyo caso todas las réplicas pertenecen a un solo grupo) o contener la misma cantidad de elementos que la cantidad de réplicas. Por ejemplo,replica_groups = {0, 2}, {1, 3}
realiza una reducción entre las réplicas0
y2
, y1
y3
.channel_id
se usa para la comunicación entre módulos: solo las operacionesall-reduce
con el mismochannel_id
pueden comunicarse entre sí.
La forma de salida es la misma que la de entrada. Por ejemplo, si hay dos réplicas y el operando tiene el valor [1.0, 2.5]
y [3.0, 5.25]
, respectivamente, en las dos réplicas, el valor de salida de este cálculo de op y suma será [4.0, 7.75]
en ambas réplicas. Si la entrada es una tupla, la salida también es una tupla.
Para calcular el resultado de AllReduce
, es necesario tener una entrada de cada réplica, por lo que, si una réplica ejecuta un nodo AllReduce
más veces que otra, la réplica anterior esperará por siempre. Dado que todas las réplicas ejecutan el mismo programa, no hay muchas formas de que eso suceda, pero es posible que la condición de un bucle while dependa de los datos de entrada, y los datos ingresados hagan que el bucle while se itere más veces en una réplica que en otra.
AllToAll
Consulta también XlaBuilder::AllToAll
.
AllToAll es una operación colectiva que envía datos de todos los núcleos a todos los núcleos. Consta de dos fases:
- La fase de dispersión. En cada núcleo, el operando se divide en una cantidad de bloques
split_count
a lo largo delsplit_dimensions
, y los bloques se dispersan en todos los núcleos, p.ej., el bloque i se envía al núcleo i. - La fase de recopilación. Cada núcleo concatena los bloques recibidos a lo largo de
concat_dimension
.
Los núcleos que participan se pueden configurar de la siguiente manera:
replica_groups
: Cada ReplicaGroup contiene una lista de los IDs de réplica que participan en el procesamiento (el ID de réplica de la réplica actual se puede recuperar conReplicaId
). AllToAll se aplicará dentro de los subgrupos en el orden especificado. Por ejemplo,replica_groups = { {1,2,3}, {4,5,0} }
significa que se aplicará un AllToAll dentro de las réplicas{1, 2, 3}
y en la fase de recopilación, y los bloques recibidos se concatenarán en el mismo orden de 1, 2, 3. Luego, se aplicará otro AllToAll en las réplicas 4, 5 y 0, y el orden de concatenación también es 4, 5, 0. Sireplica_groups
está vacío, todas las réplicas pertenecen a un grupo, en el orden de concatenación de su apariencia.
Requisitos previos:
- El tamaño de la dimensión del operando en
split_dimension
es divisible porsplit_count
. - La forma del operando no es una tupla.
AllToAll(operand, split_dimension, concat_dimension, split_count,
replica_groups)
Argumentos | Tipo | Semántica |
---|---|---|
operand |
XlaOp |
array de entrada de dimensión n |
split_dimension
|
int64
|
Un valor en el intervalo [0,
n) que nombra la dimensión en la que se divide el operando. |
concat_dimension
|
int64
|
Un valor en el intervalo [0,
n) que nombra la dimensión a través de la cual se concatenan los bloques divididos |
split_count
|
int64
|
El número de núcleos que
participan en esta operación. Si replica_groups está vacío, esta debería ser la cantidad de réplicas; de lo contrario, debería ser igual a la cantidad de réplicas en cada grupo. |
replica_groups
|
Vector de ReplicaGroup
|
Cada grupo contiene una lista de IDs de réplica. |
A continuación, se muestra un ejemplo de Alltoall.
XlaBuilder b("alltoall");
auto x = Parameter(&b, 0, ShapeUtil::MakeShape(F32, {4, 16}), "x");
AllToAll(x, /*split_dimension=*/1, /*concat_dimension=*/0, /*split_count=*/4);
En este ejemplo, hay 4 núcleos que participan en Alltoall. En cada núcleo, el operando se divide en 4 partes a lo largo de la dimensión 0, de modo que cada parte tiene la forma f32[4,4]. Las 4 partes están dispersas en todos los núcleos. Luego, cada núcleo concatena las partes recibidas en la dimensión 1, en el orden del núcleo del 0 al 4. Por lo tanto, el resultado en cada núcleo tiene la forma f32[16,4].
BatchNormGrad
Consulta también XlaBuilder::BatchNormGrad
y el documento original de normalización por lotes para obtener una descripción detallada del algoritmo.
Calcula los gradientes de la norma por lotes.
BatchNormGrad(operand, scale, mean, variance, grad_output, epsilon, feature_index)
Argumentos | Tipo | Semántica |
---|---|---|
operand |
XlaOp |
matriz de dimensiones n que se normalizará (x). |
scale |
XlaOp |
Array de 1 dimensión (\(\gamma\)) |
mean |
XlaOp |
Array de 1 dimensión (\(\mu\)) |
variance |
XlaOp |
Array de 1 dimensión (\(\sigma^2\)) |
grad_output |
XlaOp |
Gradientes pasados a BatchNormTraining (\(\nabla y\)) |
epsilon |
float |
Valor de Épsilon (\(\epsilon\)) |
feature_index |
int64 |
Índice a la dimensión del atributo en operand |
Para cada atributo de la dimensión del atributo (feature_index
es el índice de la dimensión del componente en operand
), la operación calcula los gradientes con respecto a operand
, offset
y scale
en todas las demás dimensiones. feature_index
debe ser un índice válido para la dimensión del atributo en operand
.
Los tres gradientes se definen mediante las siguientes fórmulas (si se supone que hay un array de 4 dimensiones como operand
y que tiene el índice de dimensión de los atributos l
, el tamaño del lote m
y los tamaños espaciales w
y h
):
\[ \begin{split} c_l&= \frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \left( \nabla y_{ijkl} \frac{x_{ijkl} - \mu_l}{\sigma^2_l+\epsilon} \right) \\\\ d_l&= \frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \nabla y_{ijkl} \\\\ \nabla x_{ijkl} &= \frac{\gamma_{l} }{\sqrt{\sigma^2_{l}+\epsilon} } \left( \nabla y_{ijkl} - d_l - c_l (x_{ijkl} - \mu_{l}) \right) \\\\ \nabla \gamma_l &= \sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \left( \nabla y_{ijkl} \frac{x_{ijkl} - \mu_l}{\sqrt{\sigma^2_{l}+\epsilon} } \right) \\\\\ \nabla \beta_l &= \sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h \nabla y_{ijkl} \end{split} \]
Las entradas mean
y variance
representan valores de momentos en dimensiones por lotes y espaciales.
El tipo de salida es una tupla de tres controladores:
Salidas | Tipo | Semántica |
---|---|---|
grad_operand
|
XlaOp
|
gradiente con respecto al operand de entrada ($\nabla x$) |
grad_scale
|
XlaOp
|
gradiente con respecto al scale de entrada ($\nabla \gamma$) |
grad_offset
|
XlaOp
|
gradiente con respecto a la entrada offset ($\nabla \beta$) |
BatchNormInference
Consulta también XlaBuilder::BatchNormInference
y el documento original de normalización por lotes para obtener una descripción detallada del algoritmo.
Normaliza un array en dimensiones espaciales y por lotes.
BatchNormInference(operand, scale, offset, mean, variance, epsilon, feature_index)
Argumentos | Tipo | Semántica |
---|---|---|
operand |
XlaOp |
matriz de dimensiones n que se normalizará. |
scale |
XlaOp |
Array de 1 dimensión |
offset |
XlaOp |
Array de 1 dimensión |
mean |
XlaOp |
Array de 1 dimensión |
variance |
XlaOp |
Array de 1 dimensión |
epsilon |
float |
Valor de Épsilon |
feature_index |
int64 |
Índice a la dimensión del atributo en operand |
Para cada atributo de la dimensión del componente (feature_index
es el índice de la dimensión del componente en operand
), la operación calcula la media y la varianza de todas las demás dimensiones y usa la media y la varianza para normalizar cada elemento en operand
. feature_index
debe ser un índice válido para la dimensión del componente en operand
.
BatchNormInference
equivale a llamar a BatchNormTraining
sin procesar mean
ni variance
para cada lote. En su lugar, usa los valores mean
y variance
de entrada como valores estimados. El propósito de esta op es reducir la latencia en la inferencia, de ahí el nombre BatchNormInference
.
El resultado es un array normalizado de n dimensiones con la misma forma que operand
de entrada.
BatchNormTraining
Consulta también XlaBuilder::BatchNormTraining
y the original batch normalization paper
para obtener una descripción detallada del algoritmo.
Normaliza un array en dimensiones espaciales y por lotes.
BatchNormTraining(operand, scale, offset, epsilon, feature_index)
Argumentos | Tipo | Semántica |
---|---|---|
operand |
XlaOp |
matriz de dimensiones n que se normalizará (x). |
scale |
XlaOp |
Array de 1 dimensión (\(\gamma\)) |
offset |
XlaOp |
Array de 1 dimensión (\(\beta\)) |
epsilon |
float |
Valor de Épsilon (\(\epsilon\)) |
feature_index |
int64 |
Índice a la dimensión del atributo en operand |
Para cada atributo de la dimensión del componente (feature_index
es el índice de la dimensión del componente en operand
), la operación calcula la media y la varianza de todas las demás dimensiones y usa la media y la varianza para normalizar cada elemento en operand
. feature_index
debe ser un índice válido para la dimensión del componente en operand
.
El algoritmo sigue el siguiente procedimiento para cada lote en operand
\(x\) que contenga elementos m
con w
y h
como el tamaño de las dimensiones espaciales (suponiendo que operand
es un array de 4 dimensiones):
Calcula la media del lote \(\mu_l\) para cada atributo
l
en la dimensión del atributo: \(\mu_l=\frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h x_{ijkl}\)Calcula la varianza del lote \(\sigma^2_l\): $\sigma^2l=\frac{1}{mwh}\sum{i=1}^m\sum{j=1}^w\sum{k=1}^h (x_{ijkl} - \mu_l)^2$
Normaliza, escala y cambia: \(y_{ijkl}=\frac{\gamma_l(x_{ijkl}-\mu_l)}{\sqrt[2]{\sigma^2_l+\epsilon} }+\beta_l\)
El valor épsilon, generalmente un número pequeño, se agrega para evitar errores de división por cero.
El tipo de salida es una tupla de tres XlaOp
:
Salidas | Tipo | Semántica |
---|---|---|
output
|
XlaOp
|
array n dimensional con la misma forma que la entrada operand (y) |
batch_mean |
XlaOp |
Array de 1 dimensión (\(\mu\)) |
batch_var |
XlaOp |
Array de 1 dimensión (\(\sigma^2\)) |
batch_mean
y batch_var
son momentos calculados a través de las dimensiones espaciales y del lote mediante las fórmulas anteriores.
BitcastConvertType
Consulta también XlaBuilder::BitcastConvertType
.
Al igual que tf.bitcast
en TensorFlow, realiza una operación de transmisión de bits a nivel de elementos desde una forma de datos hasta una forma objetivo. El tamaño de entrada y salida deben coincidir, p.ej., los elementos s32
se convierten en elementos f32
a través de la rutina de transmisión de bits, y un elemento s32
se convertirá en cuatro elementos s8
. Bitcast se implementa como una conversión de bajo nivel, por lo que las máquinas con diferentes representaciones de punto flotante darán diferentes resultados.
BitcastConvertType(operand, new_element_type)
Argumentos | Tipo | Semántica |
---|---|---|
operand |
XlaOp |
matriz de tipo T con atenuación D |
new_element_type |
PrimitiveType |
tipo U |
Las dimensiones del operando y de la forma objetivo deben coincidir, además de la última dimensión, que cambiará según la proporción del tamaño primitivo antes y después de la conversión.
Los tipos de elementos de origen y destino no deben ser tuplas.
Conversión de bitcast a tipo primitivo de ancho diferente
La instrucción BitcastConvert
de HLO admite el caso en el que el tamaño del tipo de elemento de salida T'
no es igual al tamaño del elemento de entrada T
. Dado que toda la operación es conceptualmente una conversión de bits y no cambia los bytes subyacentes, la forma del elemento de salida debe cambiar. Para B = sizeof(T), B' =
sizeof(T')
, hay dos casos posibles.
Primero, cuando es B > B'
, la forma de salida obtiene una nueva dimensión más pequeña de tamaño B/B'
. Por ejemplo:
f16[10,2]{1,0} %output = f16[10,2]{1,0} bitcast-convert(f32[10]{0} %input)
La regla sigue siendo la misma para los escalares eficaces:
f16[2]{0} %output = f16[2]{0} bitcast-convert(f32[] %input)
Como alternativa, para B' > B
, la instrucción requiere que la última dimensión lógica de la forma de entrada sea igual a B'/B
, y esta dimensión se descarta durante la conversión:
f32[10]{0} %output = f32[10]{0} bitcast-convert(f16[10,2]{1,0} %input)
Ten en cuenta que las conversiones entre diferentes anchos de bits no se realizan por elementos.
Transmisión
Consulta también XlaBuilder::Broadcast
.
Agrega dimensiones a un array; para ello, duplica los datos en el array.
Broadcast(operand, broadcast_sizes)
Argumentos | Tipo | Semántica |
---|---|---|
operand |
XlaOp |
El array que se duplicará. |
broadcast_sizes |
ArraySlice<int64> |
Los tamaños de las nuevas dimensiones |
Las dimensiones nuevas se insertan a la izquierda; es decir, si broadcast_sizes
tiene valores {a0, ..., aN}
y la forma de operando tiene dimensiones {b0, ..., bM}
, entonces la forma de la salida tiene las dimensiones {a0, ..., aN, b0, ..., bM}
.
El nuevo índice de dimensiones en copias del operando; p.ej.,
output[i0, ..., iN, j0, ..., jM] = operand[j0, ..., jM]
Por ejemplo, si operand
es un f32
escalar con el valor 2.0f
y broadcast_sizes
es {2, 3}
, el resultado será un array con la forma f32[2, 3]
y todos los valores en el resultado serán 2.0f
.
BroadcastInDim
Consulta también XlaBuilder::BroadcastInDim
.
Duplica el tamaño y la clasificación de un array; para ello, duplica los datos en él.
BroadcastInDim(operand, out_dim_size, broadcast_dimensions)
Argumentos | Tipo | Semántica |
---|---|---|
operand |
XlaOp |
El array que se duplicará. |
out_dim_size |
ArraySlice<int64> |
Los tamaños de las dimensiones de la forma de destino |
broadcast_dimensions |
ArraySlice<int64> |
¿A qué dimensión de la forma objetivo corresponde cada dimensión de la forma de operando? |
Es similar a Broadcast, pero permite agregar dimensiones en cualquier lugar y expandir las dimensiones existentes con tamaño 1.
operand
se transmite a la forma que describe out_dim_size
.
broadcast_dimensions
asigna las dimensiones de operand
a las dimensiones de la forma objetivo, es decir, la dimensión i del operando se asigna a la dimensión broadcast_dimension[i] de la forma de salida. Las dimensiones de operand
deben tener un tamaño de 1 o del mismo tamaño que la dimensión de la forma de salida a la que se asignan. Las dimensiones restantes se rellenan con dimensiones de tamaño 1. Luego, la transmisión de la dimensión de degeneración se transmite a través de esas dimensiones de degeneración para alcanzar la forma de salida. La semántica se describe en detalle en la página de emisión.
Call
Consulta también XlaBuilder::Call
.
Invoca un procesamiento con los argumentos dados.
Call(computation, args...)
Argumentos | Tipo | Semántica |
---|---|---|
computation |
XlaComputation |
Cálculo de tipo T_0, T_1, ..., T_{N-1} -> S con N parámetros de tipo arbitrario |
args |
secuencia de N XlaOp |
N argumentos de tipo arbitrario |
La arity y los tipos de args
deben coincidir con los parámetros de computation
. No puede tener args
.
Colesky
Consulta también XlaBuilder::Cholesky
.
Calcula la descomposición de Colesky de un lote de matrices definidas positivas simétricas (hermitas).
Cholesky(a, lower)
Argumentos | Tipo | Semántica |
---|---|---|
a |
XlaOp |
un array de rango > 2 de un tipo complejo o de punto flotante. |
lower |
bool |
si se debe usar el triángulo superior o inferior de a . |
Si lower
es true
, calcula las matrices triangulares inferiores l
, de modo que $a = l .
l^T$. Si lower
es false
, calcula las matrices triangulares superiores u
de modo que\(a = u^T . u\).
Los datos de entrada solo se leen desde el triángulo inferior/superior de a
, según el valor de lower
. Se ignoran los valores del otro triángulo. Los datos de salida se muestran en el mismo triángulo; los valores en el otro triángulo se definen en la implementación y pueden ser de cualquier tipo.
Si el rango de a
es mayor que 2, a
se trata como un lote de matrices, en el que todas, excepto las 2 dimensiones menores, son dimensiones de lote.
Si a
no es un resultado positivo simétrico (hermitiano) definido por la implementación.
Pinzas
Consulta también XlaBuilder::Clamp
.
Sujeta un operando para que se encuentre dentro del rango entre un valor mínimo y un máximo.
Clamp(min, operand, max)
Argumentos | Tipo | Semántica |
---|---|---|
min |
XlaOp |
array de tipo T |
operand |
XlaOp |
array de tipo T |
max |
XlaOp |
array de tipo T |
Con un operando y valores mínimos y máximos, muestra el operando si está en el rango entre el mínimo y el máximo; de lo contrario, muestra el valor mínimo si el operando está por debajo de este rango o el valor máximo si está por encima de este rango. Es decir, clamp(a, x, b) = min(max(a, x), b)
.
Las tres matrices deben tener la misma forma. De manera alternativa, como una forma restringida de transmisión, min
o max
pueden ser un escalar de tipo T
.
Ejemplo con min
y max
escalares:
let operand: s32[3] = {-1, 5, 9};
let min: s32 = 0;
let max: s32 = 6;
==>
Clamp(min, operand, max) = s32[3]{0, 5, 6};
Contraer
Consulta también XlaBuilder::Collapse
y la operación tf.reshape
.
Contrae las dimensiones de un array en una sola dimensión.
Collapse(operand, dimensions)
Argumentos | Tipo | Semántica |
---|---|---|
operand |
XlaOp |
array de tipo T |
dimensions |
Vector de int64 |
en orden y consecutivo de las dimensiones de T. |
La contracción reemplaza al subconjunto determinado de las dimensiones del operando por una sola dimensión. Los argumentos de entrada son un arreglo arbitrario de tipo T y un vector de constante de tiempo de compilación de índices de dimensión. Los índices de dimensión deben estar en orden (números de dimensión bajos a altos) y un subconjunto consecutivos de las dimensiones de T. Por lo tanto, {0, 1, 2}, {0, 1} o {1, 2} son todos conjuntos de dimensiones válidos, pero {1, 0} o {0, 2} no lo son. Se reemplazan por una única dimensión nueva, en la misma posición en la secuencia de dimensiones en la que se reemplazan, por el tamaño de dimensión nuevo igual al producto de los tamaños de dimensión originales. El número de dimensión más bajo en dimensions
es la dimensión que varía más lenta (la más importante) en el nido de bucles que contrae estas dimensiones, y el número de dimensión más alto varía con mayor rapidez (el más pequeño). Consulta el operador tf.reshape
si se necesita un orden de contracción más general.
Por ejemplo, supongamos que v es un array de 24 elementos:
let v = f32[4x2x3] { { {10, 11, 12}, {15, 16, 17} },
{ {20, 21, 22}, {25, 26, 27} },
{ {30, 31, 32}, {35, 36, 37} },
{ {40, 41, 42}, {45, 46, 47} } };
// Collapse to a single dimension, leaving one dimension.
let v012 = Collapse(v, {0,1,2});
then v012 == f32[24] {10, 11, 12, 15, 16, 17,
20, 21, 22, 25, 26, 27,
30, 31, 32, 35, 36, 37,
40, 41, 42, 45, 46, 47};
// Collapse the two lower dimensions, leaving two dimensions.
let v01 = Collapse(v, {0,1});
then v01 == f32[4x6] { {10, 11, 12, 15, 16, 17},
{20, 21, 22, 25, 26, 27},
{30, 31, 32, 35, 36, 37},
{40, 41, 42, 45, 46, 47} };
// Collapse the two higher dimensions, leaving two dimensions.
let v12 = Collapse(v, {1,2});
then v12 == f32[8x3] { {10, 11, 12},
{15, 16, 17},
{20, 21, 22},
{25, 26, 27},
{30, 31, 32},
{35, 36, 37},
{40, 41, 42},
{45, 46, 47} };
CollectivePermute
Consulta también XlaBuilder::CollectivePermute
.
CollectivePermute es una operación colectiva que envía y recibe datos entre réplicas.
CollectivePermute(operand, source_target_pairs)
Argumentos | Tipo | Semántica |
---|---|---|
operand |
XlaOp |
array de entrada de dimensión n |
source_target_pairs |
Vector de <int64, int64> |
Una lista de pares (source_replica_id, target_replica_id). Para cada par, el operando se envía de la réplica de origen a la réplica de destino. |
Ten en cuenta que existen las siguientes restricciones en el source_target_pair
:
- Dos pares cualesquiera no deben tener el mismo ID de réplica de destino ni el mismo ID de réplica de origen.
- Si un ID de réplica no es un destino en ningún par, la salida de esa réplica es un tensor que consta de 0 con la misma forma que la entrada.
Concatenate
Consulta también XlaBuilder::ConcatInDim
.
Concatenate compone un array a partir de múltiples operandos de array. El array tiene la misma clasificación que cada uno de los operandos de array de entrada (que deben tener el mismo rango) y contiene los argumentos en el orden en que se especificaron.
Concatenate(operands..., dimension)
Argumentos | Tipo | Semántica |
---|---|---|
operands |
secuencia de N XlaOp |
N matrices de tipo T con dimensiones [L0, L1, ...]. Requiere N >= 1. |
dimension |
int64 |
Un valor en el intervalo [0, N) que nombra la dimensión que se concatena entre los operands . |
A excepción de dimension
, todas las dimensiones deben ser iguales. Esto se debe a que XLA no admite arrays irregulares. Además, ten en cuenta que los valores de rango 0 no se pueden concatenar (ya que es imposible nombrar la dimensión a lo largo de la cual se produce la concatenación).
Ejemplo unidimensional:
Concat({ {2, 3}, {4, 5}, {6, 7} }, 0)
>>> {2, 3, 4, 5, 6, 7}
Ejemplo bidimensional:
let a = {
{1, 2},
{3, 4},
{5, 6},
};
let b = {
{7, 8},
};
Concat({a, b}, 0)
>>> {
{1, 2},
{3, 4},
{5, 6},
{7, 8},
}
Diagrama:
Condicionales
Consulta también XlaBuilder::Conditional
.
Conditional(pred, true_operand, true_computation, false_operand,
false_computation)
Argumentos | Tipo | Semántica |
---|---|---|
pred |
XlaOp |
Escalar de tipo PRED |
true_operand |
XlaOp |
Argumento de tipo \(T_0\) |
true_computation |
XlaComputation |
XlaComputation de tipo \(T_0 \to S\) |
false_operand |
XlaOp |
Argumento de tipo \(T_1\) |
false_computation |
XlaComputation |
XlaComputation de tipo \(T_1 \to S\) |
Ejecuta true_computation
si pred
es true
, false_computation
si pred
es false
, y muestra el resultado.
true_computation
debe tener un solo argumento de tipo \(T_0\) y se invocará con true_operand
, que debe ser del mismo tipo. false_computation
debe tener un solo argumento de tipo \(T_1\) y se invocará con false_operand
, que debe ser del mismo tipo. El tipo del valor que se muestra de true_computation
y false_computation
debe ser el mismo.
Ten en cuenta que solo se ejecutará uno de los valores true_computation
y false_computation
según el valor de pred
.
Conditional(branch_index, branch_computations, branch_operands)
Argumentos | Tipo | Semántica |
---|---|---|
branch_index |
XlaOp |
Escalar de tipo S32 |
branch_computations |
secuencia de N XlaComputation |
XlaComputations de tipo \(T_0 \to S , T_1 \to S , ..., T_{N-1} \to S\) |
branch_operands |
secuencia de N XlaOp |
Argumentos de tipo \(T_0 , T_1 , ..., T_{N-1}\) |
Ejecuta branch_computations[branch_index]
y muestra el resultado. Si branch_index
es una S32
que es < 0 o >= N, entonces branch_computations[N-1]
se ejecuta como la rama predeterminada.
Cada branch_computations[b]
debe tener un solo argumento de tipo \(T_b\) y se invocará con branch_operands[b]
, que debe ser del mismo tipo. El tipo del valor mostrado de cada branch_computations[b]
debe ser el mismo.
Ten en cuenta que solo se ejecutará una de las branch_computations
según el valor de branch_index
.
Conv. (convolución)
Consulta también XlaBuilder::Conv
.
Como ConvWithGeneralPadding, pero el padding se especifica de forma abreviada como SAME o VALID. SAME rellena la entrada (lhs
) con ceros para que la salida tenga la misma forma que la entrada cuando no se tienen en cuenta los pasos. Si hay relleno VALID, simplemente significa que no hay relleno.
ConvConGeneralPadding (convolución)
Consulta también XlaBuilder::ConvWithGeneralPadding
.
Calcula una convolución del tipo usado en las redes neuronales. Aquí, una convolución puede considerarse como una ventana de n dimensiones que se mueve por un área de base de esa dimensión, y se realiza un cálculo para cada posición posible de la ventana.
Argumentos | Tipo | Semántica |
---|---|---|
lhs |
XlaOp |
Array de entradas rango n+2 |
rhs |
XlaOp |
matricial n+2 de pesos del kernel |
window_strides |
ArraySlice<int64> |
Array n-d de segmentos de kernel |
padding |
ArraySlice< pair<int64,int64>> |
Array n-d de relleno (bajo, alto) |
lhs_dilation |
ArraySlice<int64> |
Array de factor de dilatación de n-d lhs |
rhs_dilation |
ArraySlice<int64> |
Array de factor de dilatación de n-d rhs |
feature_group_count |
int64 | la cantidad de grupos de atributos |
batch_group_count |
int64 | la cantidad de grupos por lotes |
Supongamos que n es el número de dimensiones espaciales. El argumento lhs
es un array de clasificación n+2 que describe el área de la base. Esto se llama entrada, aunque,
por supuesto, la RHS también es una entrada. En una red neuronal, estas son las activaciones de entrada.
Las dimensiones n+2 están en este orden:
batch
: Cada coordenada en esta dimensión representa una entrada independiente para la que se realiza la convolución.z/depth/features
: Cada posición (y,x) en el área de la base tiene un vector asociado, que entra en esta dimensión.spatial_dims
: Describe las dimensiones espacialesn
que definen el área base por la que se mueve la ventana.
El argumento rhs
es un array de rango n+2 que describe el filtro, kernel o ventana convolucional. Las dimensiones se muestran en este orden:
output-z
: Es la dimensiónz
del resultado.input-z
: El tamaño de esta dimensión multiplicado porfeature_group_count
debe igual al tamaño de la dimensiónz
en lhs.spatial_dims
: Describe las dimensiones espacialesn
que definen la ventana n-d que se mueve por el área base.
El argumento window_strides
especifica el segmento de la ventana convolucional en las dimensiones espaciales. Por ejemplo, si el segmento de la primera dimensión espacial es 3, la ventana solo se puede colocar en las coordenadas en las que el primer índice espacial es divisible por 3.
El argumento padding
especifica la cantidad de padding cero que se aplicará al área de la base. La cantidad de padding puede ser negativa; el valor absoluto del padding negativo indica la cantidad de elementos que se quitarán de la dimensión especificada antes de realizar la convolución. padding[0]
especifica el relleno de la dimensión y
y padding[1]
especifica el relleno de la dimensión x
. Cada par tiene el padding bajo como primer elemento y el padding alto como el segundo. El padding bajo se aplica en la dirección de los índices más bajos, mientras que el padding alto se aplica en la dirección de los índices más altos. Por ejemplo, si padding[1]
es (2,3)
, habrá un padding de 2 ceros a la izquierda y 3 ceros a la derecha en la segunda dimensión espacial. El uso de relleno es equivalente a insertar esos mismos valores de cero en la entrada (lhs
) antes de realizar la convolución.
Los argumentos lhs_dilation
y rhs_dilation
especifican el factor de dilatación que se aplicará a las lhs y rhs, respectivamente, en cada dimensión espacial. Si el factor de dilatación en una dimensión espacial es d, se colocan los agujeros de d-1 de forma implícita entre cada una de las entradas de esa dimensión, lo que aumenta el tamaño del array. Los agujeros se llenan con un valor no-ops, que para la convolución significa ceros.
La dilatación de la RHS también se conoce como convolución atros. Si deseas obtener información más detallada, consulta tf.nn.atrous_conv2d
. La dilatación de las lhs también se denomina convolución transpuesta. Para obtener información detallada, consulta tf.nn.conv2d_transpose
.
El argumento feature_group_count
(valor predeterminado 1) se puede usar para las convoluciones agrupadas. feature_group_count
debe ser un divisor de la dimensión del atributo de entrada y de la de salida. Si feature_group_count
es mayor que 1, significa que, conceptualmente, la dimensión del atributo de entrada y salida y la dimensión del atributo de salida rhs
se dividen de manera uniforme en muchos grupos de feature_group_count
, cada uno de los cuales consta de una subsecuencia consecutiva de atributos. La dimensión del atributo de entrada de rhs
debe ser igual a la dimensión del atributo de entrada lhs
dividida por feature_group_count
(por lo que ya tiene el tamaño de un grupo de atributos de entrada). Los grupos i-th se usan juntos para calcular feature_group_count
para muchas convoluciones distintas. Los resultados de estas convoluciones se concatenan en la dimensión del atributo de salida.
Para la convolución en profundidad, el argumento feature_group_count
se establecería en la dimensión del atributo de entrada y el filtro se cambiaría de [filter_height, filter_width, in_channels, channel_multiplier]
a [filter_height, filter_width, 1, in_channels * channel_multiplier]
. Para obtener más información, consulta tf.nn.depthwise_conv2d
.
El argumento batch_group_count
(valor predeterminado 1) se puede usar para los filtros agrupados durante la propagación inversa. batch_group_count
debe ser un divisor del tamaño de la dimensión del lote lhs
(entrada). Si batch_group_count
es mayor que 1, significa que la dimensión del lote de salida debe ser de tamaño input batch
/ batch_group_count
. El valor de batch_group_count
debe ser un divisor del tamaño del atributo de salida.
La forma de salida tiene las siguientes dimensiones, en este orden:
batch
: El tamaño de esta dimensión multiplicado porbatch_group_count
debe ser igual al tamaño de la dimensiónbatch
en lhs.z
: Tiene el mismo tamaño queoutput-z
en el kernel (rhs
).spatial_dims
: Un valor para cada posición válida de la ventana convolucional
En la figura anterior, se muestra cómo funciona el campo batch_group_count
. En efecto, dividimos cada lote en grupos batch_group_count
y hacemos lo mismo para los atributos de salida. Luego, para cada uno de estos grupos, realizamos convoluciones en pares y concatenamos la salida junto con la dimensión del atributo de salida. La semántica operativa de todas las demás dimensiones (atributo y espacial) sigue siendo la misma.
Las posiciones válidas de la ventana convolucional se determinan por los pasos y el tamaño del área de la base después del relleno.
Para describir lo que hace una convolución, considera una convolución en 2D y elige algunas coordenadas fijas batch
, z
, y
y x
en el resultado. Entonces, (y,x)
es una posición de una esquina de la ventana dentro del área de la base (p.ej., la esquina superior izquierda, según cómo interpretes las dimensiones espaciales). Ahora tenemos una ventana de 2D, tomada de la zona de la base, en la que cada punto 2D está asociado a un vector 1d, por lo que obtenemos un cuadro de 3D. Desde el kernel convolucional, desde que corregimos la coordenada de salida z
, también tenemos un cuadro 3D. Los dos cuadros tienen las mismas dimensiones, por lo que podemos tomar la suma de los productos a nivel de elementos entre los dos cuadros (similar a un producto escalar). Ese es el valor de salida.
Ten en cuenta que si output-z
es p.ej., 5, cada posición de la ventana produce 5 valores en el resultado en la dimensión z
del resultado. Estos valores difieren en qué parte del kernel convolucional se usa: hay un cuadro de valores 3D separado para cada coordenada output-z
. Puedes pensar en ello como 5 convoluciones
independientes con un filtro diferente para cada una.
A continuación, se muestra el seudocódigo de una convolución en 2D con relleno y pasos:
for (b, oz, oy, ox) { // output coordinates
value = 0;
for (iz, ky, kx) { // kernel coordinates and input z
iy = oy*stride_y + ky - pad_low_y;
ix = ox*stride_x + kx - pad_low_x;
if ((iy, ix) inside the base area considered without padding) {
value += input(b, iz, iy, ix) * kernel(oz, iz, ky, kx);
}
}
output(b, oz, oy, ox) = value;
}
ConvertElementType
Consulta también XlaBuilder::ConvertElementType
.
Al igual que static_cast
en C++, realiza una operación de conversión a nivel de elementos de una forma de datos a una forma objetivo. Las dimensiones deben coincidir, y la conversión depende de los elementos; p.ej., los elementos s32
se convierten en elementos f32
mediante una rutina de conversión de s32
a f32
.
ConvertElementType(operand, new_element_type)
Argumentos | Tipo | Semántica |
---|---|---|
operand |
XlaOp |
matriz de tipo T con atenuación D |
new_element_type |
PrimitiveType |
tipo U |
Las dimensiones del operando y de la forma objetivo deben coincidir. Los tipos de elementos de origen y de destino no deben ser tuplas.
Una conversión como T=s32
a U=f32
realizará una rutina de conversión de valor de int a float que se normaliza, como el redondeo a par más cercano.
let a: s32[3] = {0, 1, 2};
let b: f32[3] = convert(a, f32);
then b == f32[3]{0.0, 1.0, 2.0}
CrossReplicaSum
Realiza AllReduce
con un cálculo de suma.
CustomCall
Consulta también XlaBuilder::CustomCall
.
Llama a una función proporcionada por el usuario dentro de un cálculo.
CustomCall(target_name, args..., shape)
Argumentos | Tipo | Semántica |
---|---|---|
target_name |
string |
Nombre de la función. Se emitirá una instrucción de llamada orientada a este nombre de símbolo. |
args |
secuencia de N XlaOp |
N argumentos de tipo arbitrario, que se pasarán a la función. |
shape |
Shape |
Forma del resultado de la función |
La firma de la función es la misma, independientemente de la aridad o el tipo de argumentos:
extern "C" void target_name(void* out, void** in);
Por ejemplo, si CustomCall se usa de la siguiente manera:
let x = f32[2] {1,2};
let y = f32[2x3] { {10, 20, 30}, {40, 50, 60} };
CustomCall("myfunc", {x, y}, f32[3x3])
Este es un ejemplo de una implementación de myfunc
:
extern "C" void myfunc(void* out, void** in) {
float (&x)[2] = *static_cast<float(*)[2]>(in[0]);
float (&y)[2][3] = *static_cast<float(*)[2][3]>(in[1]);
EXPECT_EQ(1, x[0]);
EXPECT_EQ(2, x[1]);
EXPECT_EQ(10, y[0][0]);
EXPECT_EQ(20, y[0][1]);
EXPECT_EQ(30, y[0][2]);
EXPECT_EQ(40, y[1][0]);
EXPECT_EQ(50, y[1][1]);
EXPECT_EQ(60, y[1][2]);
float (&z)[3][3] = *static_cast<float(*)[3][3]>(out);
z[0][0] = x[1] + y[1][0];
// ...
}
La función proporcionada por el usuario no debe tener efectos secundarios y su ejecución debe ser idempotente.
Punto
Consulta también XlaBuilder::Dot
.
Dot(lhs, rhs)
Argumentos | Tipo | Semántica |
---|---|---|
lhs |
XlaOp |
array de tipo T |
rhs |
XlaOp |
array de tipo T |
La semántica exacta de esta operación depende de las clasificaciones de los operandos:
Entrada | Resultado | Semántica |
---|---|---|
vector [n] dot vector [n] |
escalar | producto vectorial de puntos |
matriz [m x k] dot vector [k] |
vector [m] | multiplicación de vectores de matrices |
matriz [m x k] dot matriz [k x n] |
matriz [m x n] | multiplicación de matrices y matrices |
La operación realiza la suma de productos en la segunda dimensión de lhs
(o
el primero si tiene clasificación 1) y la primera dimensión de rhs
. Estas son las dimensiones "contratadas". Las dimensiones contraídas de lhs
y rhs
deben ser del mismo tamaño. En la práctica, se puede usar para realizar productos de puntos entre vectores, multiplicaciones de vectores o matrices o multiplicaciones de matrices y matrices.
DotGeneral
Consulta también XlaBuilder::DotGeneral
.
DotGeneral(lhs, rhs, dimension_numbers)
Argumentos | Tipo | Semántica |
---|---|---|
lhs |
XlaOp |
array de tipo T |
rhs |
XlaOp |
array de tipo T |
dimension_numbers |
DotDimensionNumbers |
números de dimensión de lote y contratación |
Es similar a Dot, pero permite que se especifiquen números de dimensión de lote y de contratación para lhs
y rhs
.
Campos DotDimensionNumbers | Tipo | Semántica |
---|---|---|
lhs_contracting_dimensions
|
int64 repetido | Números de dimensión de contratación lhs |
rhs_contracting_dimensions
|
int64 repetido | Números de dimensión de contratación rhs |
lhs_batch_dimensions
|
int64 repetido | lhs de números de dimensión del lote |
rhs_batch_dimensions
|
int64 repetido | rhs de números de dimensión del lote |
DotGeneral realiza la suma de productos sobre las dimensiones de contratación especificadas en dimension_numbers
.
No es necesario que los números de dimensión de contratación asociados de lhs
y rhs
sean iguales, pero deben tener los mismos tamaños de dimensión.
Ejemplo con números de dimensión de contratación:
lhs = { {1.0, 2.0, 3.0},
{4.0, 5.0, 6.0} }
rhs = { {1.0, 1.0, 1.0},
{2.0, 2.0, 2.0} }
DotDimensionNumbers dnums;
dnums.add_lhs_contracting_dimensions(1);
dnums.add_rhs_contracting_dimensions(1);
DotGeneral(lhs, rhs, dnums) -> { {6.0, 12.0},
{15.0, 30.0} }
Los números de dimensión del lote asociados de lhs
y rhs
deben tener los mismos tamaños de dimensión.
Ejemplo con números de dimensión del lote (tamaño del lote 2, matrices 2x2):
lhs = { { {1.0, 2.0},
{3.0, 4.0} },
{ {5.0, 6.0},
{7.0, 8.0} } }
rhs = { { {1.0, 0.0},
{0.0, 1.0} },
{ {1.0, 0.0},
{0.0, 1.0} } }
DotDimensionNumbers dnums;
dnums.add_lhs_contracting_dimensions(2);
dnums.add_rhs_contracting_dimensions(1);
dnums.add_lhs_batch_dimensions(0);
dnums.add_rhs_batch_dimensions(0);
DotGeneral(lhs, rhs, dnums) -> { { {1.0, 2.0},
{3.0, 4.0} },
{ {5.0, 6.0},
{7.0, 8.0} } }
Entrada | Resultado | Semántica |
---|---|---|
[b0, m, k] dot [b0, k, n] |
[b0, m, n] | matmul por lotes |
[b0, b1, m, k] dot [b0, b1, k y n] |
[b0, b1, m y n]. | matmul por lotes |
A continuación, el número de dimensión resultante comienza con la dimensión del lote, luego, la dimensión lhs
no contraída/no por lotes y, por último, rhs
la dimensión no contraído/no por lotes.
DynamicSlice
Consulta también XlaBuilder::DynamicSlice
.
DynamicSlice extrae un subarreglo del arreglo de entrada en start_indices
dinámico. El tamaño de la porción en cada dimensión se pasa en size_indices
, que especifica el punto de finalización de los intervalos de porción exclusivos en cada dimensión: [inicio, inicio + tamaño). La forma de start_indices
debe tener una clasificación == 1, con un tamaño de dimensión igual al rango de operand
.
DynamicSlice(operand, start_indices, size_indices)
Argumentos | Tipo | Semántica |
---|---|---|
operand |
XlaOp |
Array N dimensional de tipo T |
start_indices |
secuencia de N XlaOp |
Lista de N números enteros escalares que contienen los índices iniciales de la porción para cada dimensión. El valor debe ser mayor o igual que cero. |
size_indices |
ArraySlice<int64> |
Lista de N números enteros que contienen el tamaño de porción de cada dimensión. Cada valor debe ser estrictamente mayor que cero, y el valor de inicio + tamaño debe ser menor o igual que el tamaño de la dimensión para evitar ajustar el tamaño de la dimensión del módulo. |
Los índices de porción eficaces se calculan mediante la aplicación de la siguiente transformación para cada índice i
en [1, N)
antes de realizar la porción:
start_indices[i] = clamp(start_indices[i], 0, operand.dimension_size[i] - size_indices[i])
Esto garantiza que la porción extraída siempre esté dentro de los límites con respecto al arreglo del operando. Si la porción está dentro de los límites antes de que se aplique la transformación, la transformación no tiene efecto.
Ejemplo unidimensional:
let a = {0.0, 1.0, 2.0, 3.0, 4.0}
let s = {2}
DynamicSlice(a, s, {2}) produces:
{2.0, 3.0}
Ejemplo bidimensional:
let b =
{ {0.0, 1.0, 2.0},
{3.0, 4.0, 5.0},
{6.0, 7.0, 8.0},
{9.0, 10.0, 11.0} }
let s = {2, 1}
DynamicSlice(b, s, {2, 2}) produces:
{ { 7.0, 8.0},
{10.0, 11.0} }
DynamicUpdateSlice
Consulta también XlaBuilder::DynamicUpdateSlice
.
DynamicUpdateSlice genera un resultado que es el valor del arreglo de entrada operand
, con una porción update
reemplazada en start_indices
.
La forma de update
determina la forma del subarreglo del resultado que se actualiza.
La forma de start_indices
debe tener una clasificación == 1, con un tamaño de dimensión igual al rango de operand
.
DynamicUpdateSlice(operand, update, start_indices)
Argumentos | Tipo | Semántica |
---|---|---|
operand |
XlaOp |
Array N dimensional de tipo T |
update |
XlaOp |
Array de dimensión N de tipo T que contiene la actualización de porción. Cada dimensión de la forma de actualización debe ser estrictamente mayor que cero, y el valor de inicio + actualización debe ser menor o igual que el tamaño del operando de cada dimensión para evitar generar índices de actualización fuera de los límites. |
start_indices |
secuencia de N XlaOp |
Lista de N números enteros escalares que contienen los índices iniciales de la porción para cada dimensión. El valor debe ser mayor o igual que cero. |
Los índices de porción eficaces se calculan mediante la aplicación de la siguiente transformación para cada índice i
en [1, N)
antes de realizar la porción:
start_indices[i] = clamp(start_indices[i], 0, operand.dimension_size[i] - update.dimension_size[i])
Esto garantiza que la porción actualizada siempre esté dentro de los límites del arreglo de operando. Si la porción está dentro de los límites antes de que se aplique la transformación, la transformación no tiene efecto.
Ejemplo unidimensional:
let a = {0.0, 1.0, 2.0, 3.0, 4.0}
let u = {5.0, 6.0}
let s = {2}
DynamicUpdateSlice(a, u, s) produces:
{0.0, 1.0, 5.0, 6.0, 4.0}
Ejemplo bidimensional:
let b =
{ {0.0, 1.0, 2.0},
{3.0, 4.0, 5.0},
{6.0, 7.0, 8.0},
{9.0, 10.0, 11.0} }
let u =
{ {12.0, 13.0},
{14.0, 15.0},
{16.0, 17.0} }
let s = {1, 1}
DynamicUpdateSlice(b, u, s) produces:
{ {0.0, 1.0, 2.0},
{3.0, 12.0, 13.0},
{6.0, 14.0, 15.0},
{9.0, 16.0, 17.0} }
Operaciones aritméticas binarias en elementos
Consulta también XlaBuilder::Add
.
Se admite un conjunto de operaciones aritméticas binarias de elementos.
Op(lhs, rhs)
Donde Op
es Add
(suma), Sub
(resta), Mul
(multiplicación), Div
(división), Rem
(resto), Max
(máximo), Min
(mínimo), LogicalAnd
(operación lógica AND) o LogicalOr
(operación lógica OR).
Argumentos | Tipo | Semántica |
---|---|---|
lhs |
XlaOp |
operando izquierdo: array de tipo T |
rhs |
XlaOp |
operando derecho: array de tipo T |
Las formas de los argumentos deben ser similares o compatibles. Consulta la documentación de transmisión sobre lo que significa que las formas sean compatibles. El resultado de una operación tiene una forma que es el resultado de transmitir los dos arrays de entrada. En esta variante, las operaciones entre arreglos de diferentes clasificaciones no son compatibles, a menos que uno de los operandos sea un escalar.
Cuando Op
es Rem
, el signo del resultado se toma del dividendo, y el valor absoluto del resultado es siempre menor que el valor absoluto del divisor.
El desbordamiento de la división de enteros (división o resto con signo o sin firma por cero o división o resto con signo de INT_SMIN
con -1
) produce un valor definido de implementación.
Existe una variante alternativa con compatibilidad de transmisión de diferente rango para estas operaciones:
Op(lhs, rhs, broadcast_dimensions)
En el ejemplo anterior, Op
es igual al anterior. Esta variante de la operación se debe usar para operaciones aritméticas entre arrays de diferentes rangos (por ejemplo, agregar una matriz a un vector).
El operando broadcast_dimensions
adicional es una porción de números enteros que se usa para expandir el rango del operando de rango inferior hasta el rango del operando de rango superior. broadcast_dimensions
asigna las dimensiones de la forma de rango inferior a las dimensiones de la forma de rango superior. Las dimensiones sin asignar de la forma expandida se rellenan con dimensiones de tamaño uno. Luego, la transmisión de la dimensión de degeneración transmite las formas a lo largo de estas dimensiones de degeneración para igualar las formas de ambos operandos. La semántica se describe en detalle en la página de emisión.
Operaciones de comparación en elementos
Consulta también XlaBuilder::Eq
.
Se admite un conjunto de operaciones estándar de comparación binaria a nivel de elementos. Ten en cuenta que se aplica la semántica de comparación de punto flotante estándar IEEE 754 cuando se comparan tipos de punto flotante.
Op(lhs, rhs)
Donde Op
es Eq
(igual a), Ne
(no igual a), Ge
(mayor o igual que), Gt
(mayor que), Le
(menor o igual que), Lt
(menor que). Otro conjunto de operadores, EqTotalOrder, NeTotalOrder, GeTotalOrder, GtTotalOrder, LeTotalOrder y LtTotalOrder, además admiten un pedido total sobre los números de punto flotante, mediante la aplicación -NaN < -Inf < -Finite < -0 < +N < +Finite <a +N <N
Argumentos | Tipo | Semántica |
---|---|---|
lhs |
XlaOp |
operando izquierdo: array de tipo T |
rhs |
XlaOp |
operando derecho: array de tipo T |
Las formas de los argumentos deben ser similares o compatibles. Consulta la documentación de transmisión sobre lo que significa que las formas sean compatibles. El resultado de una operación tiene una forma que es el resultado de transmitir los dos arrays de entrada con el tipo de elemento PRED
. En esta variante, no se admiten las operaciones entre arreglos de diferentes clasificaciones, a menos que uno de los operandos sea un escalar.
Existe una variante alternativa con compatibilidad de transmisión de diferente rango para estas operaciones:
Op(lhs, rhs, broadcast_dimensions)
En el ejemplo anterior, Op
es igual al anterior. Esta variante de la operación se debe usar para operaciones de comparación entre arrays de diferentes clasificaciones (por ejemplo, agregar una matriz a un vector).
El operando broadcast_dimensions
adicional es una porción de números enteros que especifican las dimensiones que se usarán para transmitir los operandos. La semántica se describe en detalle en la página de emisión.
Funciones unarias de elementos
XlaBuilder admite estas funciones unarias a nivel de elementos:
Abs(operand)
Abs. de elementos x -> |x|
.
Ceil(operand)
Cielo por elementos x -> ⌈x⌉
.
Cos(operand)
Coseno por elemento x -> cos(x)
.
Exp(operand)
x -> e^x
exponencial natural por elementos.
Floor(operand)
Plano básico x -> ⌊x⌋
.
Imag(operand)
Parte imaginaria de elementos de una forma compleja (o real). x -> imag(x)
. Si el operando es un tipo de punto flotante, muestra 0.
IsFinite(operand)
Comprueba si cada elemento de operand
es finito, es decir, no es infinito positivo ni negativo, y no es NaN
. Muestra un array de valores PRED
con la misma forma que la entrada, en el que cada elemento es true
solo si el elemento de entrada correspondiente es finito.
Log(operand)
Logaritmo natural por elementos x -> ln(x)
.
LogicalNot(operand)
Lógico de elemento, no x -> !(x)
.
Logistic(operand)
Cálculo de la función logística por elementos x ->
logistic(x)
.
PopulationCount(operand)
Calcula la cantidad de bits establecidos en cada elemento de operand
.
Neg(operand)
Negación por elementos x -> -x
.
Real(operand)
Parte real por elementos de una forma compleja (o real).
x -> real(x)
. Si el operando es un tipo de punto flotante, muestra el mismo valor.
Rsqrt(operand)
Recíproco a nivel de elemento de la operación de raíz cuadrada x -> 1.0 / sqrt(x)
.
Sign(operand)
Operación de signo por elementos x -> sgn(x)
, donde
\[\text{sgn}(x) = \begin{cases} -1 & x < 0\\ -0 & x = -0\\ NaN & x = NaN\\ +0 & x = +0\\ 1 & x > 0 \end{cases}\]
Usando el operador de comparación del tipo de elemento de operand
.
Sqrt(operand)
Operación de raíz cuadrada por elementos x -> sqrt(x)
.
Cbrt(operand)
Operación de raíz cúbica por elementos x -> cbrt(x)
.
Tanh(operand)
Tangente hiperbólica con elementos x -> tanh(x)
.
Round(operand)
Redondeo por elementos, empate desde cero.
RoundNearestEven(operand)
Redondeo por elementos, se vincula al par más cercano.
Argumentos | Tipo | Semántica |
---|---|---|
operand |
XlaOp |
El operando de la función |
La función se aplica a cada elemento del array operand
, lo que da como resultado un array con la misma forma. Se permite que operand
sea un escalar (rango 0).
Fft
La operación FFT de XLA implementa las Transformaciones de Fourier inversas y futuras para entradas y salidas reales y complejas. Se admiten FFT multidimensionales en hasta 3 ejes.
Consulta también XlaBuilder::Fft
.
Argumentos | Tipo | Semántica |
---|---|---|
operand |
XlaOp |
La matriz que estamos transformando Fourier. |
fft_type |
FftType |
Consulta la tabla que se encuentra a continuación |
fft_length |
ArraySlice<int64> |
Las longitudes de dominio de tiempo de los ejes que se transforman. Esto es necesario en particular para que IRFFT redimensione el eje más interno, ya que RFFT(fft_length=[16]) tiene la misma forma de salida que RFFT(fft_length=[17]) . |
FftType |
Semántica |
---|---|
FFT |
Reenvía los FFT complejos a complejos. La forma no se modifica. |
IFFT |
FFT inverso de complejo a complejo. La forma no se modifica. |
RFFT |
Reenvía los FFT reales a complejos. La forma del eje más interno se reduce a fft_length[-1] // 2 + 1 si fft_length[-1] es un valor distinto de cero, y se omite la parte conjugada invertida de la señal transformada más allá de la frecuencia de Nyquist. |
IRFFT |
FFT inverso de real a complejo (es decir, toma complejo, muestra real). La forma del eje más interno se expande a fft_length[-1] si fft_length[-1] es un valor distinto de cero, lo que infiere la parte de la señal transformada más allá de la frecuencia de Nyquist a partir del conjugado inverso de las entradas 1 a fft_length[-1] // 2 + 1 . |
FFT multidimensional
Cuando se proporciona más de 1 fft_length
, esto equivale a aplicar una cascada de operaciones de FFT a cada uno de los ejes más internos. Ten en cuenta que, en los casos reales->complejos y complejos->reales, la transformación del eje interno se realiza primero (RFFT; el último para IRFFT), por lo que el eje más interno es el que cambia de tamaño. Por lo tanto, otras transformaciones de ejes serán
complejas.
Detalles de la implementación
Los FFT de CPU cuentan con el respaldo de TensorFFT de Eigen. GPU FFT usa cuFFT.
Gather
La operación de recopilación de XLA une varias porciones (cada una en un desplazamiento de tiempo de ejecución potencialmente diferente) de un array de entrada.
Semántica general
Consulta también XlaBuilder::Gather
.
Para obtener una descripción más intuitiva, consulta la sección "Descripción informal" que aparece a continuación.
gather(operand, start_indices, offset_dims, collapsed_slice_dims, slice_sizes, start_index_map)
Argumentos | Tipo | Semántica |
---|---|---|
operand |
XlaOp |
El array del cual recopilamos. |
start_indices |
XlaOp |
Arreglo que contiene los índices iniciales de las porciones que recopilamos. |
index_vector_dim |
int64 |
Es la dimensión de start_indices que "contiene" los índices iniciales. Obtén una descripción detallada a continuación. |
offset_dims |
ArraySlice<int64> |
Conjunto de dimensiones en la forma de salida que se desplaza en un arreglo dividido del operando. |
slice_sizes |
ArraySlice<int64> |
slice_sizes[i] son los límites de la porción en la dimensión i . |
collapsed_slice_dims |
ArraySlice<int64> |
Es el conjunto de dimensiones de cada porción que se contraen. Estas dimensiones deben ser de tamaño 1. |
start_index_map |
ArraySlice<int64> |
Un mapa en el que se describe cómo asignar índices en start_indices a índices legales en un operando. |
indices_are_sorted |
bool |
Indica si se garantiza que los índices estén ordenados por el emisor. |
Para mayor comodidad, etiquetamos las dimensiones en el array de salida que no están en offset_dims
como batch_dims
.
El resultado es un array de clasificación batch_dims.size
+ offset_dims.size
.
El valor de operand.rank
debe ser igual a la suma de offset_dims.size
y collapsed_slice_dims.size
. Además, slice_sizes.size
debe ser igual a operand.rank
.
Si index_vector_dim
es igual a start_indices.rank
, consideramos implícitamente que start_indices
tenga una dimensión 1
final (es decir, si start_indices
era de forma [6,7]
y index_vector_dim
es 2
, consideramos implícitamente que la forma de start_indices
sea [6,7,1]
).
Los límites del array de salida en la dimensión i
se calculan de la siguiente manera:
Si
i
está presente enbatch_dims
(es decir, es igual abatch_dims[k]
para algunosk
), elegimos los límites de dimensión correspondientes destart_indices.shape
y se omiteindex_vector_dim
(es decir, seleccionastart_indices.shape.dims
[k
] sik
<index_vector_dim
ystart_indices.shape.dims
[k
+1
]).Si
i
está presente enoffset_dims
(es decir, igual aoffset_dims
[k
] para algunosk
), elegiremos el límite correspondiente deslice_sizes
después de considerarcollapsed_slice_dims
(es decir, seleccionamosadjusted_slice_sizes
[k
], dondeadjusted_slice_sizes
esslice_sizes
con los límites de los índicescollapsed_slice_dims
quitados).
De manera formal, el índice de operando In
correspondiente a un índice de salida Out
determinado se calcula de la siguiente manera:
Deja que
G
= {Out
[k
] parak
enbatch_dims
}. UsaG
para dividir un vectorS
de manera queS
[i
] =start_indices
[Combine(G
,i
)], en el que Combine(A, b) inserta b en la posiciónindex_vector_dim
en A. Ten en cuenta que esto está bien definido, incluso siG
está vacío: siG
está vacío, entoncesS
=start_indices
.Crea un índice inicial,
S
in
, enoperand
medianteS
mediante la dispersión deS
constart_index_map
. Más precisamente:S
in
[start_index_map
[k
]] =S
[k
] sik
<start_index_map.size
.S
in
[_
] =0
de lo contrario.
Crea un índice
O
in
enoperand
mediante la dispersión de los índices en las dimensiones de desplazamiento enOut
de acuerdo con el conjuntocollapsed_slice_dims
. Más precisamente:O
in
[remapped_offset_dims
(k
)] =Out
[offset_dims
[k
]] sik
<offset_dims.size
(remapped_offset_dims
se define a continuación).O
in
[_
] =0
de lo contrario.
In
esO
in
+S
in
, donde + es la suma de los elementos.
remapped_offset_dims
es una función monótona con dominio [0
, offset_dims.size
) y rango [0
, operand.rank
) \ collapsed_slice_dims
. Si, p.ej., offset_dims.size
es 4
, operand.rank
es 6
y collapsed_slice_dims
es {0
, 2
}, luego, remapped_offset_dims
es {0
→1
, 1
→3
, 2
→4
, 3
→5
}.
Si indices_are_sorted
se configura como verdadero, XLA puede suponer que el usuario ordena start_indices
(en orden ascendente start_index_map
). Si no lo están, entonces se define la semántica de la implementación.
Descripción informal y ejemplos
Informalmente, cada índice Out
en el array de salida corresponde a un elemento E
en el array de operando, y se calcula de la siguiente manera:
Usamos las dimensiones del lote en
Out
para buscar un índice de inicio a partir destart_indices
.Usamos
start_index_map
para asignar el índice de inicio (cuyo tamaño puede ser menor que operand.rank) a un índice de inicio "completo" enoperand
.Dividimos de forma dinámica una porción con el tamaño
slice_sizes
usando el índice de inicio completo.Para cambiar la forma de la porción, contrae las dimensiones de
collapsed_slice_dims
. Dado que todas las dimensiones de porciones contraídas deben tener un límite de 1, este cambio de forma siempre es legal.Usamos las dimensiones de desplazamiento en
Out
para indexar en esta porción y obtener el elemento de entrada,E
, que corresponde al índice de salidaOut
.
index_vector_dim
se configura como start_indices.rank
- 1
en todos los ejemplos siguientes. Los valores más interesantes para index_vector_dim
no cambian la operación en esencia, pero hacen que la representación visual sea más engorrosa.
Para tener una intuición de cómo se combina todo lo anterior, echemos un vistazo a un ejemplo que reúne 5 porciones de forma [8,6]
de un arreglo [16,11]
. La posición de una porción en el array [16,11]
se puede representar como un vector de índice de la forma S64[2]
, por lo que el conjunto de 5 posiciones se puede representar como un array S64[5,2]
.
El comportamiento de la operación de recopilación se puede representar como una transformación de índices que toma [G
,O
0
,O
1
], un índice en la forma de salida y lo asigna a un elemento en el array de entrada de la siguiente manera:
Primero, seleccionamos un vector (X
, Y
) del array de recopilación de índices con G
.
El elemento del array de salida en el índice [G
,O
0
,O
1
] es el elemento en el array de entrada en el índice [X
+ O
0
,Y
+O
1
].
slice_sizes
es [8,6]
, que decide el rango de O0
y O1
, y esto, a su vez, decide los límites de la porción.
Esta operación de recopilación actúa como una porción dinámica por lotes con G
como la dimensión del lote.
Los índices de recopilación pueden ser multidimensionales. Por ejemplo, una versión más general del ejemplo anterior que use un array de "índices de recopilación" con la forma [4,5,2]
traduciría índices como el siguiente:
Nuevamente, esto actúa como una porción dinámica por lotes G
0
y G
1
como las dimensiones del lote. El tamaño de la porción sigue siendo [8,6]
.
La operación de recopilación en XLA generaliza la semántica informal descrita anteriormente de las siguientes maneras:
Podemos configurar qué dimensiones de la forma de salida son las de desplazamiento (las que contienen
O
0
,O
1
en el último ejemplo). Las dimensiones del lote de salida (dimensiones que contienenG
0
,G
1
en el último ejemplo) se definen como las dimensiones de salida que no son dimensiones de desplazamiento.La cantidad de dimensiones de desplazamiento de salida presentes explícitamente en la forma de salida puede ser menor que la clasificación de entrada. Estas dimensiones "faltantes", que se enumeran explícitamente como
collapsed_slice_dims
, deben tener un tamaño de porción de1
. Dado que tienen un tamaño de porción de1
, el único índice válido para ellos es0
, y si los omites, no se genera ambigüedad.La porción extraída del arreglo “Reúne índices” (
X
,Y
) en el último ejemplo puede tener menos elementos que la clasificación del arreglo de entrada, y una asignación explícita determina cómo se debe expandir el índice para que tenga la misma clasificación que la entrada.
Como ejemplo final, usamos (2) y (3) para implementar tf.gather_nd
:
G
0
y G
1
se usan para dividir un índice de inicio del arreglo de índices de recopilación como de costumbre, excepto que el índice inicial tiene un solo elemento, X
. Del mismo modo, solo hay un índice de desplazamiento de salida con el valor O
0
. Sin embargo, antes de usarse como índices en el array de entrada, estos se expanden de acuerdo con "Recopilar asignación de índices" (start_index_map
en la descripción formal) y "Asignación de desplazamiento" (remapped_offset_dims
en la descripción formal) a [X
,0
] y [0
,O
0
], respectivamente, y se agregan hasta [X
,O
0
]. En otras palabras, el índice de entrada [G
0
].G
para O
la salida [G
]. G
0
0
0
0
O
1
1
GatherIndices
tf.gather_nd
slice_sizes
para este caso es [1,11]
. De forma intuitiva, esto significa que cada índice X
en el arreglo de índices de conjunto elige una fila completa y el resultado es la concatenación de todas estas filas.
GetDimensionSize
Consulta también XlaBuilder::GetDimensionSize
.
Muestra el tamaño de la dimensión determinada del operando. El operando debe tener forma de array.
GetDimensionSize(operand, dimension)
Argumentos | Tipo | Semántica |
---|---|---|
operand |
XlaOp |
array de entrada de dimensión n |
dimension |
int64 |
Un valor en el intervalo [0, n) que especifica la dimensión |
SetDimensionSize
Consulta también XlaBuilder::SetDimensionSize
.
Establece el tamaño dinámico de la dimensión determinada de XlaOp. El operando debe tener forma de array.
SetDimensionSize(operand, size, dimension)
Argumentos | Tipo | Semántica |
---|---|---|
operand |
XlaOp |
de entrada n. |
size |
XlaOp |
int32 que representa el tamaño dinámico del entorno de ejecución |
dimension |
int64 |
Es un valor en el intervalo [0, n) que especifica la dimensión. |
Pasa el operando como resultado, con la dimensión dinámica que realiza el seguimiento del compilador.
Las operaciones de reducción descendentes ignorarán los valores con relleno.
let v: f32[10] = f32[10]{1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
let five: s32 = 5;
let six: s32 = 6;
// Setting dynamic dimension size doesn't change the upper bound of the static
// shape.
let padded_v_five: f32[10] = set_dimension_size(v, five, /*dimension=*/0);
let padded_v_six: f32[10] = set_dimension_size(v, six, /*dimension=*/0);
// sum == 1 + 2 + 3 + 4 + 5
let sum:f32[] = reduce_sum(padded_v_five);
// product == 1 * 2 * 3 * 4 * 5
let product:f32[] = reduce_product(padded_v_five);
// Changing padding size will yield different result.
// sum == 1 + 2 + 3 + 4 + 5 + 6
let sum:f32[] = reduce_sum(padded_v_six);
GetTupleElement
Consulta también XlaBuilder::GetTupleElement
.
Índices en una tupla con un valor constante de tiempo de compilación.
El valor debe ser una constante de tiempo de compilación para que la inferencia de forma pueda determinar el tipo del valor resultante.
Esto es análogo a std::get<int N>(t)
en C++. Conceptualmente:
let v: f32[10] = f32[10]{0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
let s: s32 = 5;
let t: (f32[10], s32) = tuple(v, s);
let element_1: s32 = gettupleelement(t, 1); // Inferred shape matches s32.
Consulta también tf.tuple
.
Entrada
Consulta también XlaBuilder::Infeed
.
Infeed(shape)
Argumento | Tipo | Semántica |
---|---|---|
shape |
Shape |
Es la forma de los datos leídos en la interfaz de Infeed. El campo de diseño de la forma se debe configurar para que coincida con el diseño de los datos enviados al dispositivo; de lo contrario, su comportamiento no está definido. |
Lee un solo elemento de datos desde la interfaz de transmisión de entrada implícita del dispositivo, interpreta los datos como la forma determinada y su diseño, y muestra una XlaOp
de los datos. Se permiten varias operaciones de entrada en un cálculo, pero debe haber un orden total entre ellas. Por ejemplo, dos entradas del siguiente código tienen un orden total, ya que hay una dependencia entre los bucles while.
result1 = while (condition, init = init_value) {
Infeed(shape)
}
result2 = while (condition, init = result1) {
Infeed(shape)
}
Las formas de tuplas anidadas no son compatibles. En el caso de una forma de tupla vacía, la operación de entrada es, de hecho, una no-op y continúa sin leer ningún dato de la entrada del dispositivo.
Pizca
Consulta también XlaBuilder::Iota
.
Iota(shape, iota_dimension)
Compila un literal constante en el dispositivo en lugar de una transferencia de host potencialmente grande. Crea un array con una forma especificada y contiene valores que comienzan en cero y que aumentan en uno a lo largo de la dimensión especificada. Para los tipos de punto flotante, el array producido es equivalente a ConvertElementType(Iota(...))
, en el que Iota
es de tipo integral y la conversión es al tipo de punto flotante.
Argumentos | Tipo | Semántica |
---|---|---|
shape |
Shape |
Forma del array creado por Iota() |
iota_dimension |
int64 |
Es la dimensión que se incrementa. |
Por ejemplo, Iota(s32[4, 8], 0)
muestra
[[0, 0, 0, 0, 0, 0, 0, 0 ],
[1, 1, 1, 1, 1, 1, 1, 1 ],
[2, 2, 2, 2, 2, 2, 2, 2 ],
[3, 3, 3, 3, 3, 3, 3, 3 ]]
Devoluciones por Iota(s32[4, 8], 1)
[[0, 1, 2, 3, 4, 5, 6, 7 ],
[0, 1, 2, 3, 4, 5, 6, 7 ],
[0, 1, 2, 3, 4, 5, 6, 7 ],
[0, 1, 2, 3, 4, 5, 6, 7 ]]
Asignar
Consulta también XlaBuilder::Map
.
Map(operands..., computation)
Argumentos | Tipo | Semántica |
---|---|---|
operands |
secuencia de N XlaOp |
N matrices de tipos T0..T{N-1} |
computation |
XlaComputation |
Cálculo de tipo T_0, T_1, .., T_{N + M -1} -> S con N parámetros de tipo T y M de tipo arbitrario |
dimensions |
Array int64 |
array de dimensiones del mapa |
Aplica una función escalar sobre los arrays operands
dados, lo que produce un array con las mismas dimensiones en el que cada elemento es el resultado de la función asignada aplicada a los elementos correspondientes de los arrays de entrada.
La función asignada es un cálculo arbitrario con la restricción de que tiene N entradas del tipo escalar T
y una sola salida con el tipo S
. La salida tiene las mismas dimensiones que los operandos, excepto que el tipo de elemento T se reemplaza por S.
Por ejemplo: Map(op1, op2, op3, computation, par1)
asigna elem_out <-
computation(elem1, elem2, elem3, par1)
a cada índice (multidimensional) en los arreglos de entrada para producir el arreglo de salida.
OptimizationBarrier
Bloquea cualquier pase de optimización para que no desplace los cálculos a través de la barrera.
Garantiza que todas las entradas se evalúen antes que los operadores que dependen de las salidas de la barrera.
Almohadilla
Consulta también XlaBuilder::Pad
.
Pad(operand, padding_value, padding_config)
Argumentos | Tipo | Semántica |
---|---|---|
operand |
XlaOp |
array de tipo T |
padding_value |
XlaOp |
escalar de tipo T para completar el padding agregado. |
padding_config |
PaddingConfig |
la cantidad de padding en ambos bordes (bajo, alto) y entre los elementos de cada dimensión |
Expande el array operand
dado mediante padding alrededor de él y entre sus elementos con el padding_value
especificado. padding_config
especifica la cantidad de padding del borde y padding interior de cada dimensión.
PaddingConfig
es un campo repetido de PaddingConfigDimension
, que contiene tres campos para cada dimensión: edge_padding_low
, edge_padding_high
y interior_padding
.
edge_padding_low
y edge_padding_high
especifican la cantidad de padding que se agrega en el extremo inferior (junto al índice 0) y en el extremo alto (junto al índice más alto) de cada dimensión, respectivamente. La cantidad de padding del borde puede ser negativa; el valor absoluto del padding negativo indica la cantidad de elementos que se quitarán de la dimensión especificada.
interior_padding
especifica la cantidad de padding que se agrega entre dos elementos cualesquiera de cada dimensión. Este valor no puede ser negativo. El relleno interior se produce de forma lógica antes del relleno del borde, por lo que, en el caso del relleno del borde negativo, los elementos se quitan del operando con relleno interno.
Esta operación es una no-op si todos los pares de padding de bordes son (0, 0) y los valores de padding interior son todos 0. En la siguiente figura, se muestran ejemplos de diferentes valores edge_padding
y interior_padding
para un array bidimensional.
Recv
Consulta también XlaBuilder::Recv
.
Recv(shape, channel_handle)
Argumentos | Tipo | Semántica |
---|---|---|
shape |
Shape |
forma de los datos para recibir |
channel_handle |
ChannelHandle |
identificador único para cada par de envío/recepción |
Recibe datos de la forma determinada de una instrucción Send
en otro cálculo que comparte el mismo controlador de canal. Muestra una XlaOp para los datos recibidos.
La API de cliente de la operación Recv
representa la comunicación síncrona.
Sin embargo, la instrucción se divide de forma interna en 2 instrucciones HLO (Recv
y RecvDone
) para habilitar las transferencias de datos asíncronas. Consulta también HloInstruction::CreateRecv
y HloInstruction::CreateRecvDone
.
Recv(const Shape& shape, int64 channel_id)
Asigna los recursos necesarios para recibir datos de una instrucción Send
con el mismo channel_id. Muestra un contexto para los recursos asignados, que se usa con una instrucción RecvDone
siguiente a fin de esperar que se complete la transferencia de datos. El contexto es una tupla de {ángulo de recepción (forma), identificador de solicitud (U32)} y solo se puede usar mediante una instrucción RecvDone
.
RecvDone(HloInstruction context)
En un contexto creado por una instrucción Recv
, espera a que se complete la transferencia de datos y muestra los datos recibidos.
Reducir
Consulta también XlaBuilder::Reduce
.
Aplica una función de reducción a uno o más arrays en paralelo.
Reduce(operands..., init_values..., computation, dimensions)
Argumentos | Tipo | Semántica |
---|---|---|
operands |
Secuencia de N XlaOp |
N arrays de tipos T_0, ..., T_{N-1} . |
init_values |
Secuencia de N XlaOp |
N escalares de tipos T_0, ..., T_{N-1} . |
computation |
XlaComputation |
el cálculo del tipo T_0, ..., T_{N-1}, T_0, ..., T_{N-1} -> Collate(T_0, ..., T_{N-1}) . |
dimensions |
Array int64 |
y sin ordenar de las dimensiones que se reducirán. |
Aquí:
- N debe ser mayor o igual que 1.
- El cálculo debe ser asociativo “aproximadamente” (consulta a continuación).
- Todos los arrays de entrada deben tener las mismas dimensiones.
- Todos los valores iniciales deben formar una identidad inferior a
computation
. - Si es
N = 1
,Collate(T)
esT
. - Si es
N > 1
,Collate(T_0, ..., T_{N-1})
es una tupla de elementosN
de tipoT
.
Esta operación reduce una o más dimensiones de cada array de entrada en escalares.
La clasificación de cada arreglo que se muestra es rank(operand) - len(dimensions)
. El resultado de la operación es Collate(Q_0, ..., Q_N)
, en el que Q_i
es un array de tipo T_i
, cuyas dimensiones se describen a continuación.
Diferentes backends pueden volver a asociar el cálculo de la reducción. Esto puede generar diferencias numéricas, ya que algunas funciones de reducción, como la suma, no son asociativas con los números de punto flotante. Sin embargo, si el rango de los datos es limitado, la adición de punto flotante está lo suficientemente cerca de ser asociativa para los usos más prácticos.
Ejemplos
Cuando se realiza una reducción en una dimensión en un solo array 1D con valores [10, 11,
12, 13]
, con la función de reducción f
(esto es computation
), eso se podría
calcular como
f(10, f(11, f(12, f(init_value, 13)))
pero también existen muchas otras posibilidades, p.ej.
f(init_value, f(f(10, f(init_value, 11)), f(f(init_value, 12), f(init_value, 13))))
El siguiente es un ejemplo aproximado de un seudocódigo de cómo se puede implementar la reducción, con la sumatoria como el cálculo de la reducción con un valor inicial de 0.
result_shape <- remove all dims in dimensions from operand_shape
# Iterate over all elements in result_shape. The number of r's here is equal
# to the rank of the result
for r0 in range(result_shape[0]), r1 in range(result_shape[1]), ...:
# Initialize this result element
result[r0, r1...] <- 0
# Iterate over all the reduction dimensions
for d0 in range(dimensions[0]), d1 in range(dimensions[1]), ...:
# Increment the result element with the value of the operand's element.
# The index of the operand's element is constructed from all ri's and di's
# in the right order (by construction ri's and di's together index over the
# whole operand shape).
result[r0, r1...] += operand[ri... di]
Este es un ejemplo de cómo reducir un array 2D (matriz). La forma tiene rango 2, dimensión 0 de tamaño 2 y dimensión 1 de tamaño 3:
Resultados de reducir las dimensiones 0 o 1 con una función "add":
Ten en cuenta que ambos resultados de reducción son arrays 1D. En el diagrama, se muestra uno como columna y otro como fila solo para mayor conveniencia visual.
Para un ejemplo más complejo, aquí hay un array 3D. Su clasificación es 3, dimensión 0 de tamaño 4, dimensión 1 de tamaño 2 y dimensión 2 de tamaño 3. Para simplificar, los valores 1 a 6 se replican en la dimensión 0.
Al igual que en el ejemplo en 2D, podemos reducir solo una dimensión. Si reducimos la dimensión 0, por ejemplo, obtenemos un array de rango 2 en el que todos los valores de la dimensión 0 se pliegan en un escalar:
| 4 8 12 |
| 16 20 24 |
Si reducimos la dimensión 2, también obtenemos un array de rango 2 en el que todos los valores de la dimensión 2 se plegaron en un escalar:
| 6 15 |
| 6 15 |
| 6 15 |
| 6 15 |
Ten en cuenta que el orden relativo entre las dimensiones restantes de la entrada se conserva en la salida, pero es posible que a algunas dimensiones se les asignen números nuevos (debido a que cambia la clasificación).
También podemos reducir varias dimensiones. De reducción de las dimensiones 0 y 1, se produce un array 1D [20, 28, 36]
.
Si reduces el array 3D en todas sus dimensiones, se produce el 84
escalar.
Reducción variable
Cuando es N > 1
, la aplicación de la función reducir es un poco más compleja, ya que se aplica simultáneamente a todas las entradas. Los operandos se suministran al cálculo en el siguiente orden:
- Ejecutando valor reducido para el primer operando
- ...
- Ejecutando valor reducido para el operando enésimo
- Valor de entrada para el primer operando
- ...
- Valor de entrada para el operando enésimo
Por ejemplo, considera la siguiente función de reducción, que se puede usar para calcular el máximo y el argmax de un array de 1D en paralelo:
f: (Float, Int, Float, Int) -> Float, Int
f(max, argmax, value, index):
if value >= max:
return (value, index)
else:
return (max, argmax)
Para los arrays de entrada 1D V = Float[N], K = Int[N]
y los valores init I_V = Float, I_K = Int
, el resultado f_(N-1)
de la reducción en la única dimensión de entrada es equivalente a la siguiente aplicación recursiva:
f_0 = f(I_V, I_K, V_0, K_0)
f_1 = f(f_0.first, f_0.second, V_1, K_1)
...
f_(N-1) = f(f_(N-2).first, f_(N-2).second, V_(N-1), K_(N-1))
Si aplicas esta reducción a un array de valores y un array de índices secuenciales (es decir, iota), se iterarán en conjunto los arrays y se mostrará una tupla que contenga el valor máximo y el índice coincidente.
ReducePrecision
Consulta también XlaBuilder::ReducePrecision
.
Modela el efecto de convertir valores de punto flotante a un formato de menor precisión (como IEEE-FP16) y de vuelta al formato original. La cantidad de bits de exponentes y mantisas en el formato de menor precisión se puede especificar de manera arbitraria, aunque es posible que no todos los tamaños de bits sean compatibles con todas las implementaciones de hardware.
ReducePrecision(operand, mantissa_bits, exponent_bits)
Argumentos | Tipo | Semántica |
---|---|---|
operand |
XlaOp |
array de tipo de punto flotante T . |
exponent_bits |
int32 |
cantidad de bits exponentes en un formato de menor precisión |
mantissa_bits |
int32 |
cantidad de bits de mantisa en un formato de menor precisión |
El resultado es un array de tipo T
. Los valores de entrada se redondean al valor más cercano representable con la cantidad determinada de bits de mantisa (mediante la semántica de "empates a pares"), y cualquier valor que exceda el rango especificado por la cantidad de bits exponentes se restringen al infinito positivo o negativo. Se conservan los valores NaN
, aunque se pueden convertir en valores NaN
canónicos.
El formato de menor precisión debe tener al menos un bit del exponente (para distinguir un valor cero de un valor infinito, ya que ambos tienen una mantisa cero) y debe tener un número no negativo de bits de mantisa. La cantidad de bits de exponente o mantisa puede exceder el valor correspondiente para el tipo T
. La parte correspondiente de la conversión es simplemente una no-op.
ReduceScatter
Consulta también XlaBuilder::ReduceScatter
.
ReduceScatter es una operación colectiva que realiza un AllReduce y, luego, dispersa el resultado dividiéndolo en bloques shard_count
a lo largo de scatter_dimension
y la réplica i
del grupo de réplicas recibe el fragmento ith
.
ReduceScatter(operand, computation, scatter_dim, shard_count,
replica_group_ids, channel_id)
Argumentos | Tipo | Semántica |
---|---|---|
operand |
XlaOp |
Es un array o una tupla no vacía de arrays que se reducirán en las réplicas. |
computation |
XlaComputation |
Cálculo de reducción |
scatter_dimension |
int64 |
Dimensión que se debe dispersar. |
shard_count |
int64 |
Cantidad de bloques para dividir scatter_dimension |
replica_groups |
vector de vectores de int64 |
Grupos entre los que se realizan las reducciones |
channel_id |
int64 opcional |
ID de canal opcional para la comunicación entre módulos |
- Cuando
operand
es una tupla de arrays, se realiza la reducción de dispersión en cada elemento de la tupla. replica_groups
es una lista de grupos de réplicas entre los que se realiza la reducción (el ID de réplica de la réplica actual se puede recuperar conReplicaId
). El orden de las réplicas en cada grupo determina el orden en el que se dispersará el resultado de All-reduce.replica_groups
debe estar vacío (en ese caso, todas las réplicas pertenecen a un solo grupo) o contener la misma cantidad de elementos que la cantidad de réplicas. Cuando hay más de un grupo de réplicas, todos deben tener el mismo tamaño. Por ejemplo,replica_groups = {0, 2}, {1, 3}
realiza una reducción entre las réplicas0
y2
, y1
y3
, y, luego, dispersa el resultado.shard_count
es el tamaño de cada grupo de réplicas. Esto lo necesitamos en los casos en que los objetosreplica_groups
están vacíos. Sireplica_groups
no está vacío,shard_count
debe ser igual al tamaño de cada grupo de réplicas.channel_id
se usa para la comunicación entre módulos: solo las operacionesreduce-scatter
con el mismochannel_id
pueden comunicarse entre sí.
La forma de salida es la de entrada, con scatter_dimension
hecho shard_count
veces más pequeño. Por ejemplo, si hay dos réplicas y el operando tiene el valor [1.0, 2.25]
y [3.0, 5.25]
respectivamente en ambas réplicas, el valor de salida de esta operación en la que scatter_dim
es 0
será [4.0]
para la primera réplica y [7.5]
para la segunda.
ReduceWindow
Consulta también XlaBuilder::ReduceWindow
.
Aplica una función de reducción a todos los elementos en cada ventana de una secuencia de N arrays multidimensionales, lo que produce una única o una tupla de N arrays multidimensionales como resultado. Cada arreglo de salida tiene la misma cantidad de elementos que la cantidad de posiciones válidas de la ventana. Una capa de reducción se puede expresar como ReduceWindow
. Al igual que Reduce
, al objeto computation
aplicado siempre se le pasa el init_values
en el lado izquierdo.
ReduceWindow(operands..., init_values..., computation, window_dimensions,
window_strides, padding)
Argumentos | Tipo | Semántica |
---|---|---|
operands |
N XlaOps |
Es una secuencia de N arrays multidimensionales de tipos T_0,..., T_{N-1} , cada uno de los cuales representa el área base en la que se coloca la ventana. |
init_values |
N XlaOps |
Los N valores iniciales para la reducción, uno para cada uno de los N operandos. Consulta Reducir para obtener más detalles. |
computation |
XlaComputation |
Función de reducción de tipo T_0, ..., T_{N-1}, T_0, ..., T_{N-1} -> Collate(T_0, ..., T_{N-1}) , para aplicar a elementos en cada ventana de todos los operandos de entrada. |
window_dimensions |
ArraySlice<int64> |
array de números enteros para los valores de las dimensiones de ventana |
window_strides |
ArraySlice<int64> |
Array de números enteros para los valores de segmento de ventana |
base_dilations |
ArraySlice<int64> |
array de números enteros para valores de dilatación base |
window_dilations |
ArraySlice<int64> |
array de números enteros para valores de dilatación de ventanas |
padding |
Padding |
el tipo de padding para la ventana (Padding::kSame, que se rellena para tener la misma forma de salida que la entrada si el stride es 1, o Padding::kValid, que no usa relleno y “detiene” la ventana cuando ya no encaja) |
Aquí:
- N debe ser mayor o igual que 1.
- Todos los arrays de entrada deben tener las mismas dimensiones.
- Si es
N = 1
,Collate(T)
esT
. - Si es
N > 1
,Collate(T_0, ..., T_{N-1})
es una tupla de elementosN
de tipo(T0,...T{N-1})
.
Debajo del código y la figura, se muestra un ejemplo del uso de ReduceWindow
. La entrada es una matriz de tamaño [4x6], y tanto window_dimensions como window_stride_dimensions son [2x3].
// Create a computation for the reduction (maximum).
XlaComputation max;
{
XlaBuilder builder(client_, "max");
auto y = builder.Parameter(0, ShapeUtil::MakeShape(F32, {}), "y");
auto x = builder.Parameter(1, ShapeUtil::MakeShape(F32, {}), "x");
builder.Max(y, x);
max = builder.Build().value();
}
// Create a ReduceWindow computation with the max reduction computation.
XlaBuilder builder(client_, "reduce_window_2x3");
auto shape = ShapeUtil::MakeShape(F32, {4, 6});
auto input = builder.Parameter(0, shape, "input");
builder.ReduceWindow(
input,
/*init_val=*/builder.ConstantLiteral(LiteralUtil::MinValue(F32)),
*max,
/*window_dimensions=*/{2, 3},
/*window_stride_dimensions=*/{2, 3},
Padding::kValid);
El segmento de 1 en una dimensión especifica que la posición de una ventana en la dimensión se encuentra a 1 elemento de distancia de la ventana adyacente. Para especificar que ninguna ventana se superpone entre sí, window_stride_dimensions debe ser igual a window_dimensions. En la siguiente figura, se ilustra el uso de dos valores de segmento diferentes. El relleno se aplica a cada dimensión de la entrada y los cálculos son los mismos que si la entrada hubiera llegado con las dimensiones que tenía después del relleno.
Para un ejemplo de padding no trivial, considera calcular el mínimo de ventana reducida (el valor inicial es MAX_FLOAT
) con la dimensión 3
y el segmento 2
sobre el array de entrada [10000, 1000, 100, 10, 1]
. El padding kValid
calcula los mínimos en dos ventanas válidas: [10000, 1000, 100]
y [100, 10, 1]
, lo que genera el resultado [100, 1]
. Con el padding kSame
, primero se rellena el array para que la forma después de la ventana de reducción sea igual que la entrada para el segmento uno. Para ello, agrega elementos iniciales en ambos lados y obtén [MAX_VALUE, 10000, 1000, 100, 10, 1,
MAX_VALUE]
. La ejecución de la función "reduce-window" en el array con padding opera en tres ventanas: [MAX_VALUE, 10000, 1000]
, [1000, 100, 10]
, [10, 1, MAX_VALUE]
y produce [1000, 10, 1]
.
El orden de evaluación de la función de reducción es arbitrario y puede ser no determinista. Por lo tanto, la función de reducción no debe ser demasiado sensible a la reasociación. Consulta el debate sobre la asociatividad en el contexto de Reduce
para obtener más detalles.
ReplicaId
Consulta también XlaBuilder::ReplicaId
.
Muestra el ID único (escalar U32) de la réplica.
ReplicaId()
El ID único de cada réplica es un número entero sin firmar en el intervalo [0, N)
, en el que N
es la cantidad de réplicas. Dado que todas las réplicas ejecutan el mismo programa, una llamada ReplicaId()
en el programa mostrará un valor diferente en cada réplica.
Reshape
Consulta también XlaBuilder::Reshape
y la operación Collapse
.
Cambia la forma de las dimensiones de un array a una configuración nueva.
Reshape(operand, new_sizes)
Reshape(operand, dimensions, new_sizes)
Argumentos | Tipo | Semántica |
---|---|---|
operand |
XlaOp |
array de tipo T |
dimensions |
Vector de int64 |
orden en el que se contraen las dimensiones |
new_sizes |
Vector de int64 |
vector de tamaños de nuevas dimensiones |
Conceptualmente, primero, redimensionar compacta un array en un vector unidimensional de valores de datos y, luego, define mejor este vector en una nueva forma. Los argumentos de entrada son un arreglo arbitrario de tipo T, un vector de constante de tiempo de compilación de índices de dimensión y un vector de constante de tiempo de compilación de tamaños de dimensión para el resultado.
Los valores en el vector dimension
, si se proporcionan, deben ser una permutación de todas las dimensiones de T; el valor predeterminado si no se proporciona es {0, ..., rank - 1}
. El orden de las dimensiones en dimensions
va desde la dimensión más lenta que varía (la más importante) hasta la que varía más rápida (la más secundaria) en el anidamiento de bucles que contrae el array de entrada en una sola dimensión. El vector new_sizes
determina el tamaño del array de salida. El valor del índice 0 de new_sizes
es el tamaño de la dimensión 0, el valor del índice 1 es el tamaño de la dimensión 1, y así sucesivamente. El producto de las dimensiones new_size
debe ser igual al producto de los tamaños de dimensión del operando. Cuando defines mejor el array contraído en el array multidimensional definido por new_sizes
, las dimensiones en new_sizes
se ordenan desde la variación más lenta (la más importante) hasta la más rápida (la más secundaria).
Por ejemplo, supongamos que v es un array de 24 elementos:
let v = f32[4x2x3] { { {10, 11, 12}, {15, 16, 17} },
{ {20, 21, 22}, {25, 26, 27} },
{ {30, 31, 32}, {35, 36, 37} },
{ {40, 41, 42}, {45, 46, 47} } };
In-order collapse:
let v012_24 = Reshape(v, {0,1,2}, {24});
then v012_24 == f32[24] {10, 11, 12, 15, 16, 17, 20, 21, 22, 25, 26, 27,
30, 31, 32, 35, 36, 37, 40, 41, 42, 45, 46, 47};
let v012_83 = Reshape(v, {0,1,2}, {8,3});
then v012_83 == f32[8x3] { {10, 11, 12}, {15, 16, 17},
{20, 21, 22}, {25, 26, 27},
{30, 31, 32}, {35, 36, 37},
{40, 41, 42}, {45, 46, 47} };
Out-of-order collapse:
let v021_24 = Reshape(v, {1,2,0}, {24});
then v012_24 == f32[24] {10, 20, 30, 40, 11, 21, 31, 41, 12, 22, 32, 42,
15, 25, 35, 45, 16, 26, 36, 46, 17, 27, 37, 47};
let v021_83 = Reshape(v, {1,2,0}, {8,3});
then v021_83 == f32[8x3] { {10, 20, 30}, {40, 11, 21},
{31, 41, 12}, {22, 32, 42},
{15, 25, 35}, {45, 16, 26},
{36, 46, 17}, {27, 37, 47} };
let v021_262 = Reshape(v, {1,2,0}, {2,6,2});
then v021_262 == f32[2x6x2] { { {10, 20}, {30, 40},
{11, 21}, {31, 41},
{12, 22}, {32, 42} },
{ {15, 25}, {35, 45},
{16, 26}, {36, 46},
{17, 27}, {37, 47} } };
Como caso especial, remodelar puede transformar un array de un solo elemento en un escalar y viceversa. Por ejemplo,
Reshape(f32[1x1] { {5} }, {0,1}, {}) == 5;
Reshape(5, {}, {1,1}) == f32[1x1] { {5} };
Rev (reversión)
Consulta también XlaBuilder::Rev
.
Rev(operand, dimensions)
Argumentos | Tipo | Semántica |
---|---|---|
operand |
XlaOp |
array de tipo T |
dimensions |
ArraySlice<int64> |
dimensiones para revertir |
Invierte el orden de los elementos en el array operand
a lo largo del dimensions
especificado, lo que genera un array de salida de la misma forma. Cada elemento del array de operandos en un índice multidimensional se almacena en el array de salida en un índice transformado. Para transformar el índice multidimensional, se invierte el índice de cada dimensión que se revertirá (es decir, si una dimensión de tamaño N es una de las dimensiones invertidas, su índice i se transforma en N - 1 - i).
Un uso de la operación Rev
es revertir el array de pesos de convolución junto con las dos dimensiones de ventana durante el cálculo de gradientes en redes neuronales.
RngNormal
Consulta también XlaBuilder::RngNormal
.
Construye un resultado de una forma dada con números al azar generados siguiendo la \(N(\mu, \sigma)\) distribución normal. Los parámetros \(\mu\) , \(\sigma\)y la forma de salida deben tener un tipo de elemento de punto flotante. Además, los parámetros deben tener valores escalares.
RngNormal(mu, sigma, shape)
Argumentos | Tipo | Semántica |
---|---|---|
mu |
XlaOp |
Escalar de tipo T que especifica la media de los números generados |
sigma |
XlaOp |
Escalar de tipo T que especifica la desviación estándar del modelo generado |
shape |
Shape |
Forma de salida de tipo T |
RngUniform
Consulta también XlaBuilder::RngUniform
.
Construye un resultado de una forma determinada con números aleatorios generados a partir de la distribución uniforme durante el intervalo \([a,b)\). Los parámetros y el tipo de elemento de salida deben ser de tipo booleano, integral o de punto flotante, y los tipos deben ser coherentes. En la actualidad, los backends de CPU y GPU solo admiten F64, F32, F16, BF16, S64, U64, S32 y U32. Además, los parámetros deben tener valores escalares. Si \(b <= a\) el resultado está definido por la implementación.
RngUniform(a, b, shape)
Argumentos | Tipo | Semántica |
---|---|---|
a |
XlaOp |
Escalar de tipo T que especifica el límite de intervalo más bajo |
b |
XlaOp |
Escalar de tipo T que especifica el límite superior del intervalo |
shape |
Shape |
Forma de salida de tipo T |
RngBitGenerator
Genera un resultado con una forma determinada rellena de bits aleatorios uniformes mediante el algoritmo especificado (o el valor predeterminado del backend) y muestra un estado actualizado (con la misma forma que el estado inicial) y los datos aleatorios generados.
El estado inicial es el estado inicial de la generación actual de números aleatorios. Este, la forma requerida y los valores válidos dependen del algoritmo utilizado.
Se garantiza que el resultado será una función determinista del estado inicial, pero no se garantiza que lo sea entre los backends y las diferentes versiones del compilador.
RngBitGenerator(algorithm, key, shape)
Argumentos | Tipo | Semántica |
---|---|---|
algorithm |
RandomAlgorithm |
Se usará el algoritmo PRNG. |
initial_state |
XlaOp |
Estado inicial del algoritmo PRNG |
shape |
Shape |
Forma de la salida para los datos generados. |
Valores disponibles para algorithm
:
rng_default
: Algoritmo específico del backend con requisitos de forma específicos del backend.rng_three_fry
: Es el algoritmo de PRNG basado en contador de ThreeFry. La formainitial_state
esu64[2]
con valores arbitrarios. Salmon et al. SC 2011. Números aleatorios paralelos: tan fácil como 1, 2, 3rng_philox
: Es el algoritmo Philox para generar números aleatorios en paralelo. La formainitial_state
esu64[3]
con valores arbitrarios. Salmon et al. SC 2011. Números aleatorios paralelos: tan fácil como 1, 2, 3
Dispersión
La operación de dispersión de XLA genera una secuencia de resultados que son valores del arreglo de entrada operands
, con varias porciones (en los índices especificados por scatter_indices
) actualizadas con la secuencia de valores en updates
mediante update_computation
.
Consulta también XlaBuilder::Scatter
.
scatter(operands..., scatter_indices, updates..., update_computation,
index_vector_dim, update_window_dims, inserted_window_dims,
scatter_dims_to_operand_dims)
Argumentos | Tipo | Semántica |
---|---|---|
operands |
Secuencia de N XlaOp |
N arrays de tipos T_0, ..., T_N en los que se dispersará. |
scatter_indices |
XlaOp |
Arreglo que contiene los índices iniciales de las porciones en las que se deben dispersar. |
updates |
Secuencia de N XlaOp |
N arrays de tipos T_0, ..., T_N . updates[i] contiene los valores que deben usarse para la dispersión operands[i] . |
update_computation |
XlaComputation |
Cálculo que se usa para combinar los valores existentes en el array de entrada y las actualizaciones durante la dispersión. Este cálculo debe ser del tipo T_0, ..., T_N, T_0, ..., T_N -> Collate(T_0, ..., T_N) . |
index_vector_dim |
int64 |
Es la dimensión en scatter_indices que contiene los índices iniciales. |
update_window_dims |
ArraySlice<int64> |
Es el conjunto de dimensiones en forma updates que son dimensiones de ventana. |
inserted_window_dims |
ArraySlice<int64> |
Es el conjunto de dimensiones de la ventana que se deben insertar con la forma updates . |
scatter_dims_to_operand_dims |
ArraySlice<int64> |
Se asignan dimensiones de los índices de dispersión al espacio de índices de operando. Este array se interpreta como la asignación de i a scatter_dims_to_operand_dims[i] . Tienen que ser uno a uno y totales. |
indices_are_sorted |
bool |
Indica si se garantiza que los índices estén ordenados por el emisor. |
unique_indices |
bool |
Indica si se garantiza que los índices sean únicos para el emisor. |
Aquí:
- N debe ser mayor o igual que 1.
operands
[0
], ... yoperands
[N-1
] deben tener las mismas dimensiones.updates
[0
], ... yupdates
[N-1
] deben tener las mismas dimensiones.- Si es
N = 1
,Collate(T)
esT
. - Si es
N > 1
,Collate(T_0, ..., T_N)
es una tupla de elementosN
de tipoT
.
Si index_vector_dim
es igual a scatter_indices.rank
, de manera implícita consideramos que scatter_indices
tiene una dimensión 1
final.
Definimos update_scatter_dims
de tipo ArraySlice<int64>
como el conjunto de dimensiones en la forma updates
que no están en update_window_dims
, en orden ascendente.
Los argumentos de dispersión deben seguir estas restricciones:
Cada array
updates
debe tener la clasificaciónupdate_window_dims.size + scatter_indices.rank - 1
.Los límites de la dimensión
i
en cada arrayupdates
deben cumplir con lo siguiente:- Si
i
está presente enupdate_window_dims
(es decir, igual aupdate_window_dims
[k
] para algunak
), el límite de la dimensióni
enupdates
no debe exceder el límite correspondiente deoperand
después de contemplarinserted_window_dims
(es decir,adjusted_window_bounds
[k
], dondeadjusted_window_bounds
contiene los límites deoperand
sin los límites de los índicesinserted_window_dims
). - Si
i
está presente enupdate_scatter_dims
(es decir, igual aupdate_scatter_dims
[k
] para algunosk
), el límite de la dimensióni
enupdates
debe ser igual al límite correspondiente descatter_indices
y se omiteindex_vector_dim
(es decir,scatter_indices.shape.dims
[k
], sik
<index_vector_dim
yscatter_indices.shape.dims
[k+1
]).
- Si
update_window_dims
debe estar en orden ascendente, no debe tener ningún número de dimensión que se repita y estar en el rango[0, updates.rank)
.inserted_window_dims
debe estar en orden ascendente, no debe tener ningún número de dimensión que se repita y estar en el rango[0, operand.rank)
.operand.rank
debe ser igual a la suma deupdate_window_dims.size
yinserted_window_dims.size
.scatter_dims_to_operand_dims.size
debe ser igual ascatter_indices.shape.dims
[index_vector_dim
] y sus valores deben estar en el rango[0, operand.rank)
.
Para un índice determinado U
en cada array updates
, el índice I
correspondiente en el array operands
correspondiente al que se debe aplicar esta actualización se calcula de la siguiente manera:
- Deja que
G
= {U
[k
] parak
enupdate_scatter_dims
}. UsaG
para buscar un vector de índiceS
en el arrayscatter_indices
de modo queS
[i
] =scatter_indices
[Combine(G
,i
)] donde Combine(A, b) inserte b en las posicionesindex_vector_dim
en A. - Crea un índice
S
in
enoperand
conS
mediante la dispersiónS
con el mapascatter_dims_to_operand_dims
. Más formalmente:S
in
[scatter_dims_to_operand_dims
[k
]] =S
[k
] sik
<scatter_dims_to_operand_dims.size
.S
in
[_
] =0
de lo contrario.
- Crea un índice
W
in
en cada arrayoperands
. Para ello, dispersa los índices enupdate_window_dims
enU
de acuerdo coninserted_window_dims
. Más formalmente:W
in
[window_dims_to_operand_dims
(k
)] =U
[k
] sik
está enupdate_window_dims
, dondewindow_dims_to_operand_dims
es la función monótona con el dominio [0
,update_window_dims.size
) y el rango [0
,operand.rank
) \inserted_window_dims
. (Por ejemplo, siupdate_window_dims.size
es4
,operand.rank
es6
, yinserted_window_dims
es {0
,2
},window_dims_to_operand_dims
es {0
→1
,1
→3
,2
→4
,3
→5
}).W
in
[_
] =0
de lo contrario.
I
esW
in
+S
in
, donde + es la suma de los elementos.
En resumen, la operación de dispersión puede definirse de la siguiente manera.
- Inicializa
output
conoperands
, es decir, para todos los índicesJ
, para todos los índicesO
del arrayoperands
[J
]:
output
[J
][O
] =operands
[J
][O
] - Para cada índice
U
del arrayupdates
[J
] y el índice correspondienteO
en el arrayoperand
[J
], siO
es un índice válido paraoutput
:
(output
[0
][O
], ...,output
[N-1
][O
]) =update_computation
(output
[0
][O
], ..., ,output
[N-1
][O
],updates
[0
][U
], ...,updates
[N-1
][U
])
El orden en el que se aplican las actualizaciones no es determinista. Por lo tanto, cuando varios índices de updates
hacen referencia al mismo índice en operands
, el valor correspondiente en output
no será determinista.
Ten en cuenta que el primer parámetro que se pasa a update_computation
siempre será el valor actual del array output
y el segundo siempre será el valor del array updates
. Esto es importante en particular para los casos en los que update_computation
no es conmutativo.
Si indices_are_sorted
se configura como verdadero, XLA puede suponer que el usuario ordena start_indices
(en orden ascendente start_index_map
). Si no lo están, entonces se define la semántica de la implementación.
Si unique_indices
se establece como verdadero, XLA puede suponer que todos los elementos dispersos son únicos. Por lo tanto, XLA podría usar operaciones no atómicas. Si unique_indices
se configura como verdadero y los índices que se dispersan no son únicos, se define la implementación de la semántica.
informalmente, la op de dispersión puede verse como una inversión de la op de recopilación, es decir, la op de dispersión actualiza los elementos en la entrada que son extraídos por la op de recopilación correspondiente.
Para obtener una descripción informal detallada y ejemplos, consulta la sección "Descripción informal" en Gather
.
Seleccionar
Consulta también XlaBuilder::Select
.
Construye un arreglo de salida a partir de elementos de dos arreglos de entrada, en función de los valores de un arreglo de predicado.
Select(pred, on_true, on_false)
Argumentos | Tipo | Semántica |
---|---|---|
pred |
XlaOp |
array de tipo PRED |
on_true |
XlaOp |
array de tipo T |
on_false |
XlaOp |
array de tipo T |
Los arrays on_true
y on_false
deben tener la misma forma. Esta también es la forma del array de salida. El array pred
debe tener la misma dimensionalidad que on_true
y on_false
, con el tipo de elemento PRED
.
Para cada elemento P
de pred
, el elemento correspondiente del array de salida se toma de on_true
si el valor de P
es true
y de on_false
si el valor de P
es false
. Como una forma restringida de transmisión, pred
puede ser un escalar de tipo PRED
. En este caso, el array de salida se toma por completo de on_true
si pred
es true
y de on_false
si pred
es false
.
Ejemplo con pred
no escalar:
let pred: PRED[4] = {true, false, false, true};
let v1: s32[4] = {1, 2, 3, 4};
let v2: s32[4] = {100, 200, 300, 400};
==>
Select(pred, v1, v2) = s32[4]{1, 200, 300, 4};
Ejemplo con pred
escalar:
let pred: PRED = true;
let v1: s32[4] = {1, 2, 3, 4};
let v2: s32[4] = {100, 200, 300, 400};
==>
Select(pred, v1, v2) = s32[4]{1, 2, 3, 4};
Se admiten las selecciones entre tuplas. Para este propósito, se considera que las tuplas
son tipos escalares. Si on_true
y on_false
son tuplas (que deben tener la misma forma), pred
debe ser un escalar de tipo PRED
.
SelectAndScatter
Consulta también XlaBuilder::SelectAndScatter
.
Esta operación puede considerarse como una operación compuesta que primero calcula ReduceWindow
en el array operand
para seleccionar un elemento de cada ventana y, luego, dispersa el array source
en los índices de los elementos seleccionados para construir un array de salida con la misma forma que el array de operando. La función binaria select
se usa para seleccionar un elemento de cada ventana aplicándolo en cada ventana, y se llama con la propiedad de que el vector de índice del primer parámetro es lexicográficamente menor que el vector de índice del segundo parámetro. La función select
muestra true
si se selecciona el primer parámetro y false
si se selecciona el segundo, y la función debe contener transitividad (es decir, si select(a, b)
y select(b, c)
son true
, entonces select(a, c)
también es true
) para que el elemento seleccionado no dependa del orden de los elementos recorridos en una ventana determinada.
La función scatter
se aplica en cada índice seleccionado en el array de salida. Requiere dos parámetros escalares:
- Valor actual en el índice seleccionado en el array de salida
- El valor de dispersión de
source
que se aplica al índice seleccionado
Combina los dos parámetros y muestra un valor escalar que se usa para actualizar el valor en el índice seleccionado en el arreglo de salida. Inicialmente, todos los índices del arreglo de salida se configuran en init_value
.
El array de salida tiene la misma forma que el array operand
, y el array source
debe tener la misma forma que el resultado de aplicar una operación ReduceWindow
en el array operand
. Se puede usar SelectAndScatter
para propagar de forma inversa los valores de gradiente de una capa de reducción en una red neuronal.
SelectAndScatter(operand, select, window_dimensions, window_strides,
padding, source, init_value, scatter)
Argumentos | Tipo | Semántica |
---|---|---|
operand |
XlaOp |
Array de tipo T sobre el que se deslizan las ventanas |
select |
XlaComputation |
Cálculo binario de tipo T, T -> PRED , para aplicar a todos los elementos en cada ventana. Muestra true si se selecciona el primer parámetro y false si se selecciona el segundo parámetro. |
window_dimensions |
ArraySlice<int64> |
array de números enteros para los valores de las dimensiones de ventana |
window_strides |
ArraySlice<int64> |
Array de números enteros para los valores de segmento de ventana |
padding |
Padding |
Tipo de padding para la ventana (Padding::kSame o Padding::kValid) |
source |
XlaOp |
un array de tipo T con los valores que se van a dispersar. |
init_value |
XlaOp |
valor escalar de tipo T para el valor inicial del array de salida |
scatter |
XlaComputation |
Cálculo binario de tipo T, T -> T , para aplicar cada elemento de fuente de dispersión con su elemento de destino |
En la siguiente figura, se muestran ejemplos del uso de SelectAndScatter
con la función select
que calcula el valor máximo entre sus parámetros. Ten en cuenta que cuando las ventanas se superponen, como en la figura (2) de abajo, diferentes ventanas pueden seleccionar un índice del array operand
varias veces. En la figura, las dos ventanas superiores (azul y rojo) seleccionan el elemento de valor 9 y la función scatter
de suma binaria produce el elemento de resultado del valor 8 (2 + 6).
El orden de evaluación de la función scatter
es arbitrario y puede ser no determinista. Por lo tanto, la función scatter
no debe ser muy sensible a la reasociación. Consulta el debate sobre la asociatividad en el contexto de Reduce
para obtener más detalles.
Enviar
Consulta también XlaBuilder::Send
.
Send(operand, channel_handle)
Argumentos | Tipo | Semántica |
---|---|---|
operand |
XlaOp |
datos para enviar (matriz de tipo T) |
channel_handle |
ChannelHandle |
identificador único para cada par de envío/recepción |
Envía los datos del operando dados a una instrucción Recv
en otro cálculo que comparte el mismo controlador de canal. No muestra ningún dato.
Al igual que la operación Recv
, la API de cliente de la operación Send
representa la comunicación síncrona y se divide de manera interna en 2 instrucciones HLO (Send
y SendDone
) para habilitar las transferencias de datos asíncronas. Consulta también HloInstruction::CreateSend
y HloInstruction::CreateSendDone
.
Send(HloInstruction operand, int64 channel_id)
Inicia una transferencia asíncrona del operando a los recursos asignados por la instrucción Recv
con el mismo ID de canal. Muestra un contexto, que se usa en una instrucción SendDone
siguiente para esperar que se complete la transferencia de datos. El contexto es una tupla de {operando (forma), identificador de solicitud (U32)} y solo se puede usar con una instrucción SendDone
.
SendDone(HloInstruction context)
En un contexto creado por una instrucción Send
, espera a que se complete la transferencia de datos. Esta instrucción no muestra ningún dato.
Programación de instrucciones del canal
El orden de ejecución de las 4 instrucciones para cada canal (Recv
, RecvDone
, Send
, SendDone
) es el siguiente.
Recv
ocurre antes delSend
Send
ocurre antes delRecvDone
Recv
ocurre antes delRecvDone
Send
ocurre antes delSendDone
Cuando los compiladores de backend generan una programación lineal para cada procesamiento que se comunica a través de las instrucciones del canal, no debe haber ciclos a través de los cálculos. Por ejemplo, los siguientes programas generan interbloqueos.
Porción
Consulta también XlaBuilder::Slice
.
La segmentación extrae un subarray del array de entrada. El subarreglo es del mismo rango que la entrada y contiene los valores dentro de un cuadro delimitador dentro del array de entrada, en el que las dimensiones y los índices del cuadro de límite se proporcionan como argumentos a la operación de división.
Slice(operand, start_indices, limit_indices, strides)
Argumentos | Tipo | Semántica |
---|---|---|
operand |
XlaOp |
Array N dimensional de tipo T |
start_indices |
ArraySlice<int64> |
Lista de N números enteros que contienen los índices iniciales de la porción para cada dimensión. Los valores deben ser mayores o iguales que cero. |
limit_indices |
ArraySlice<int64> |
Lista de N números enteros que contienen los índices finales (exclusivos) para la porción de cada dimensión. Cada valor debe ser mayor o igual que el valor de start_indices respectivo de la dimensión, así como menor o igual que el tamaño de la dimensión. |
strides |
ArraySlice<int64> |
Lista de N números enteros que deciden el segmento de entrada de la porción. La porción elige cada elemento strides[d] en la dimensión d . |
Ejemplo unidimensional:
let a = {0.0, 1.0, 2.0, 3.0, 4.0}
Slice(a, {2}, {4}) produces:
{2.0, 3.0}
Ejemplo bidimensional:
let b =
{ {0.0, 1.0, 2.0},
{3.0, 4.0, 5.0},
{6.0, 7.0, 8.0},
{9.0, 10.0, 11.0} }
Slice(b, {2, 1}, {4, 3}) produces:
{ { 7.0, 8.0},
{10.0, 11.0} }
Ordenar
Consulta también XlaBuilder::Sort
.
Sort(operands, comparator, dimension, is_stable)
Argumentos | Tipo | Semántica |
---|---|---|
operands |
ArraySlice<XlaOp> |
Los operandos que se van a ordenar. |
comparator |
XlaComputation |
El cálculo del comparador que se usará. |
dimension |
int64 |
La dimensión en la que se ordenará. |
is_stable |
bool |
Si se debe usar la ordenación estable. |
Si solo se proporciona un operando:
Si el operando es un tensor de rango-1 (un array), el resultado es un array ordenado. Si deseas ordenar el array en orden ascendente, el comparador debe realizar una comparación de menor que. De manera formal, después de ordenar el array, se mantiene en todas las posiciones de índice
i, j
coni < j
quecomparator(value[i], value[j]) = comparator(value[j], value[i]) = false
ocomparator(value[i], value[j]) = true
.Si el operando tiene una clasificación más alta, se ordena según la dimensión proporcionada. Por ejemplo, para un tensor de rango-2 (una matriz), un valor de dimensión de
0
ordenará cada columna de forma independiente, y un valor de dimensión de1
ordenará cada fila de forma independiente. Si no se proporciona un número de dimensión, se elige la última de forma predeterminada. Para la dimensión que está ordenada, se aplica el mismo orden de clasificación que en el caso de rango-1.
Si se proporcionan operandos n > 1
:
Todos los operandos
n
deben ser tensores con las mismas dimensiones. Los tipos de elementos de los tensores pueden ser diferentes.Todos los operandos se ordenan juntos, no de manera individual. Conceptualmente, los operandos se tratan como una tupla. Cuando se verifica si los elementos de cada operando en las posiciones de índice
i
yj
deben intercambiarse, se llama al comparador con los parámetros escalares2 * n
, en los que el parámetro2 * k
corresponde al valor en la posicióni
del operandok-th
y el parámetro2 * k + 1
corresponde al valor en la posiciónj
del operandok-th
. Por lo general, el comparador compararía los parámetros2 * k
y2 * k + 1
entre sí y, posiblemente, usaría otros pares de parámetros como rompientes de empate.El resultado es una tupla que consta de los operandos en orden (en la dimensión proporcionada, como se muestra arriba). El operando
i-th
de la tupla corresponde al operandoi-th
de Sort.
Por ejemplo, si hay tres operandos operand0 = [3, 1]
, operand1 = [42, 50]
, operand2 = [-3.0, 1.1]
, y el comparador solo compara los valores de operand0
con menor que, el resultado del orden es la tupla ([1, 3], [50, 42], [1.1, -3.0])
.
Si is_stable
se configura como verdadero, se garantiza que el orden sea estable, es decir, si hay elementos que el comparador considera iguales, se conserva el orden relativo de los valores iguales. Dos elementos e1
y e2
son iguales solo si comparator(e1, e2) = comparator(e2, e1) = false
. Según la configuración predeterminada, is_stable
se configura como falso.
Transposición
Consulta también la operación tf.reshape
.
Transpose(operand)
Argumentos | Tipo | Semántica |
---|---|---|
operand |
XlaOp |
El operando que se va a transponer. |
permutation |
ArraySlice<int64> |
Cómo permutar las dimensiones |
Permuta las dimensiones del operando con la permutación determinada, por lo tanto, ∀ i . 0 ≤ i < rank ⇒ input_dimensions[permutation[i]] = output_dimensions[i]
.
Es lo mismo que Reshape(operando, permutación, Permute(permutación, operando.shape.dimensions)).
TriangularSolve
Consulta también XlaBuilder::TriangularSolve
.
Resuelve sistemas de ecuaciones lineales con matrices de coeficiente triangular inferior o superior mediante sustitución hacia adelante o atrás. Cuando se transmite a través de dimensiones principales, esta rutina resuelve uno de los sistemas de matriz op(a) * x =
b
o x * op(a) = b
para la variable x
, dadas a
y b
, en las que op(a)
es op(a) = a
, op(a) = Transpose(a)
o op(a) = Conj(Transpose(a))
.
TriangularSolve(a, b, left_side, lower, unit_diagonal, transpose_a)
Argumentos | Tipo | Semántica |
---|---|---|
a |
XlaOp |
un array de rango > 2 de un tipo de punto flotante o complejo con la forma [..., M, M] . |
b |
XlaOp |
un array de rango > 2 del mismo tipo con la forma [..., M, K] si left_side es verdadero, [..., K, M] de lo contrario |
left_side |
bool |
indica si se debe resolver un sistema con el formato op(a) * x = b (true ) o x * op(a) = b (false ). |
lower |
bool |
si se debe usar el triángulo superior o inferior de a . |
unit_diagonal |
bool |
si es true , se supone que los elementos diagonales de a son 1 y no se accede a ellos. |
transpose_a |
Transpose |
si se debe usar a tal como está, se transpone o toma su transposición conjugada. |
Los datos de entrada solo se leen desde el triángulo inferior/superior de a
, según el valor de lower
. Se ignoran los valores del otro triángulo. Los datos de salida se muestran en el mismo triángulo; los valores en el otro triángulo se definen en la implementación y pueden ser de cualquier tipo.
Si la clasificación de a
y b
es mayor que 2, se tratan como lotes de matrices y todas, excepto las 2 dimensiones menores, son dimensiones de lote. a
y b
deben tener las mismas dimensiones de lote.
Tupla
Consulta también XlaBuilder::Tuple
.
Es una tupla que contiene un número variable de controladores de datos, cada uno de los cuales tiene su propia forma.
Esto es análogo a std::tuple
en C++. Conceptualmente:
let v: f32[10] = f32[10]{0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
let s: s32 = 5;
let t: (f32[10], s32) = tuple(v, s);
Se puede deconstruir (se puede acceder a las tuplas) a través de la operación GetTupleElement
.
Mientras que
Consulta también XlaBuilder::While
.
While(condition, body, init)
Argumentos | Tipo | Semántica |
---|---|---|
condition |
XlaComputation |
XlaComputation de tipo T -> PRED , que define la condición de finalización del bucle. |
body |
XlaComputation |
XlaComputation de tipo T -> T que define el cuerpo del bucle. |
init |
T |
Valor inicial para el parámetro de condition y body . |
Ejecuta secuencialmente body
hasta que condition
falla. Esto es similar a un bucle while típico en muchos otros idiomas, excepto por las diferencias y restricciones que se indican a continuación.
- Un nodo
While
muestra un valor de tipoT
, que es el resultado de la última ejecución debody
. - La forma del tipo
T
se determina de forma estática y debe ser la misma en todas las iteraciones.
Los parámetros T de los cálculos se inicializan con el valor init
en la primera iteración y se actualizan automáticamente al resultado nuevo desde body
en cada iteración posterior.
Un caso de uso principal del nodo While
es implementar la ejecución repetida del entrenamiento en redes neuronales. A continuación, se muestra un seudocódigo simplificado con un gráfico que representa el cálculo. Puedes encontrar el código en while_test.cc
.
El tipo T
en este ejemplo es un Tuple
que consta de un int32
para el recuento de iteraciones y un vector[10]
para el acumulador. Para 1,000 iteraciones, el bucle
agrega un vector constante al acumulador.
// Pseudocode for the computation.
init = {0, zero_vector[10]} // Tuple of int32 and float[10].
result = init;
while (result(0) < 1000) {
iteration = result(0) + 1;
new_vector = result(1) + constant_vector[10];
result = {iteration, new_vector};
}