以下では、XlaBuilder
インターフェースで定義されているオペレーションのセマンティクスについて説明します。通常、これらのオペレーションは、xla_data.proto
の RPC インターフェースで定義されたオペレーションに 1 対 1 でマッピングされます。
命名法に関する注意: XLA で一般化されるデータ型は、ユニフォーム型(32 ビット浮動小数点など)の要素を保持する N 次元配列です。このドキュメントでは、任意の次元の配列を表すために array を使用します。便宜上、特殊なケースにはより具体的でわかりやすい名前が付けられています。たとえば、ベクトルは 1 次元配列、行列は 2 次元配列です。
AfterAll
XlaBuilder::AfterAll
もご覧ください。
AfterAll は、可変長数のトークンを受け取り、単一のトークンを生成します。トークンはプリミティブ型で、副作用オペレーションの間でスレッド化して順序を適用できます。AfterAll
は、set オペレーションの後にオペレーションを順序付けるためのトークンの結合として使用できます。
AfterAll(operands)
引数 | タイプ | セマンティクス |
---|---|---|
operands |
XlaOp |
トークンの可変長 |
AllGather
XlaBuilder::AllGather
もご覧ください。
レプリカ間で連結を実行します。
AllGather(operand, all_gather_dim, shard_count, replica_group_ids,
channel_id)
引数 | タイプ | セマンティクス |
---|---|---|
operand
|
XlaOp
|
レプリカ間で連結するための配列 |
all_gather_dim |
int64 |
連結ディメンション |
replica_groups
|
int64 のベクトルのベクトル |
連結が行われるグループ間で |
channel_id
|
省略可 int64
|
モジュール間通信用のオプションのチャネル ID |
replica_groups
は、連結が行われるレプリカ グループのリストです(現在のレプリカのレプリカ ID はReplicaId
を使用して取得できます)。各グループ内のレプリカの順序によって、入力が結果に配置される順序が決まります。replica_groups
は空にするか(この場合、すべてのレプリカが0
からN - 1
の順に 1 つのグループに属します)、またはレプリカの数と同じ数の要素が含まれている必要があります。たとえば、replica_groups = {0, 2}, {1, 3}
はレプリカ0
と2
と、1
と3
を連結します。shard_count
は、各レプリカ グループのサイズです。これは、replica_groups
が空の場合に必要になります。channel_id
はモジュール間通信に使用されます。互いに通信できるのは、同じchannel_id
を持つall-gather
オペレーションのみです。
出力シェイプは、all_gather_dim
が shard_count
倍になった入力シェイプです。たとえば、2 つのレプリカがあり、2 つのレプリカでオペランドの値が [1.0, 2.5]
と [3.0, 5.25]
の場合、all_gather_dim
が 0
であるこの演算の出力値は両方のレプリカで [1.0, 2.5, 3.0,
5.25]
になります。
AllReduce
XlaBuilder::AllReduce
もご覧ください。
レプリカ間でカスタム計算を実行します。
AllReduce(operand, computation, replica_group_ids, channel_id)
引数 | タイプ | セマンティクス |
---|---|---|
operand
|
XlaOp
|
レプリカ全体で削減する配列または空でない配列のタプル |
computation |
XlaComputation |
削減の計算 |
replica_groups
|
int64 のベクトルのベクトル |
リダクションが行われるグループと |
channel_id
|
省略可 int64
|
モジュール間通信用のオプションのチャネル ID |
operand
が配列のタプルの場合、タプルの各要素に対して all-reduce が実行されます。replica_groups
は、リダクションが行われるレプリカ グループのリストです(現在のレプリカのレプリカ ID はReplicaId
を使用して取得できます)。replica_groups
は空にするか(この場合、すべてのレプリカが 1 つのグループに属している)、またはレプリカの数と同じ数の要素が含まれている必要があります。たとえば、replica_groups = {0, 2}, {1, 3}
はレプリカ0
と2
と、1
と3
の間でリダクションを行います。channel_id
はモジュール間通信に使用されます。互いに通信できるのは、同じchannel_id
を持つall-reduce
オペレーションのみです。
出力シェイプは入力シェイプと同じです。たとえば、2 つのレプリカがあり、2 つのレプリカでオペランドがそれぞれ [1.0, 2.5]
と [3.0, 5.25]
の値を持つ場合、この演算と合計計算からの出力値は両方のレプリカで [4.0, 7.75]
になります。入力がタプルの場合、出力もタプルです。
AllReduce
の結果を計算するには、各レプリカから 1 つの入力が必要です。したがって、あるレプリカが別のレプリカよりも AllReduce
ノードを何度も実行すると、前のレプリカは永久に待機します。レプリカはすべて同じプログラムを実行しているため、それが発生する方法はあまりありませんが、while ループの条件がインフィードのデータに依存しており、インフィードされたデータが原因で while ループが 1 つのレプリカで他のレプリカよりも繰り返し反復される場合に発生する可能性があります。
AllToAll
XlaBuilder::AllToAll
もご覧ください。
AllToAll は、すべてのコアからすべてのコアにデータを送信する集合オペレーションです。次の 2 つのフェーズがあります。
- 分散フェーズ。各コアで、オペランドは
split_dimensions
に沿ってsplit_count
個のブロックに分割され、ブロックはすべてのコアに分散されます。たとえば、i 番目のブロックは i 番目のコアに送信されます。 - 収集フェーズ。各コアは、受信したブロックを
concat_dimension
に沿って連結します。
参加するコアは、次のように構成できます。
replica_groups
: 各 ReplicaGroup には、計算に参加するレプリカ ID のリストが含まれます(現在のレプリカのレプリカ ID は、ReplicaId
を使用して取得できます)。AllToAll は指定された順序でサブグループ内に適用されます。たとえば、replica_groups = { {1,2,3}, {4,5,0} }
は、AllToAll がレプリカ{1, 2, 3}
内と収集フェーズで適用され、受信したブロックが 1、2、3 の同じ順序で連結されることを意味します。次に、レプリカ 4、5、0 内に別の AllToAll が適用され、連結順序も 4、5、0 になります。replica_groups
が空の場合、すべてのレプリカが、外観上の連結順に 1 つのグループに属します。
前提条件:
split_dimension
のオペランドのディメンション サイズはsplit_count
で割り切れます。- オペランドの形状がタプルではない。
AllToAll(operand, split_dimension, concat_dimension, split_count,
replica_groups)
引数 | タイプ | セマンティクス |
---|---|---|
operand |
XlaOp |
n 次元の入力配列 |
split_dimension
|
int64
|
オペランドを分割するディメンションの名前を指定する、[0,
n) 区間の値 |
concat_dimension
|
int64
|
分割されたブロックが連結されるディメンションの名前を示す間隔 [0,
n) の値 |
split_count
|
int64
|
このオペレーションに参加するコアの数。replica_groups が空の場合はレプリカの数と同じにする必要があります。それ以外の場合は、各グループのレプリカの数と同じにする必要があります。 |
replica_groups
|
ReplicaGroup ベクトル
|
各グループにはレプリカ ID のリストが含まれます。 |
以下は 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);
この例では、Alltoall に参加しているコアが 4 つあります。各コアで、オペランドは次元 0 に沿って 4 つの部分に分割されるため、各部分の形状は f32[4,4] になります。4 つの部分はすべてのコアに分散されています。次に、各コアは、受け取ったパーツをディメンション 1 に沿って、コア 0 ~ 4 の順序で連結します。したがって、各コアの出力の形状は f32[16,4] です。
BatchNormGrad
アルゴリズムの詳細については、XlaBuilder::BatchNormGrad
と元のバッチ正規化に関する論文もご覧ください。
バッチノルムの勾配を計算します。
BatchNormGrad(operand, scale, mean, variance, grad_output, epsilon, feature_index)
引数 | タイプ | セマンティクス |
---|---|---|
operand |
XlaOp |
正規化する n 次元配列(x) |
scale |
XlaOp |
1 次元配列(\(\gamma\)) |
mean |
XlaOp |
1 次元配列(\(\mu\)) |
variance |
XlaOp |
1 次元配列(\(\sigma^2\)) |
grad_output |
XlaOp |
BatchNormTraining (\(\nabla y\))に渡される勾配 |
epsilon |
float |
イプシロン値(\(\epsilon\)) |
feature_index |
int64 |
operand の特徴ディメンションのインデックス |
特徴ディメンションの特徴(feature_index
は operand
の特徴ディメンションのインデックス)ごとに、他のすべてのディメンションの operand
、offset
、scale
に関する勾配を計算します。feature_index
は、operand
内の特徴ディメンションの有効なインデックスである必要があります。
3 つの勾配は、次の式で定義されます(4 次元配列を operand
、特徴ディメンションのインデックス l
、バッチサイズ m
、空間サイズ w
と 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} \]
入力 mean
と variance
は、バッチ ディメンションと空間ディメンション全体でのモーメント値を表します。
出力タイプは、次の 3 つのハンドルのタプルです。
出力 | タイプ | セマンティクス |
---|---|---|
grad_operand
|
XlaOp
|
入力 operand に関する勾配($\nabla x$) |
grad_scale
|
XlaOp
|
入力 scale に関する勾配($\nabla \gamma$) |
grad_offset
|
XlaOp
|
入力 offset ($\nabla \beta$)に関する勾配 |
BatchNormInference
アルゴリズムの詳細については、XlaBuilder::BatchNormInference
と元のバッチ正規化に関する論文もご覧ください。
バッチ ディメンションと空間ディメンション全体で配列を正規化します。
BatchNormInference(operand, scale, offset, mean, variance, epsilon, feature_index)
引数 | タイプ | セマンティクス |
---|---|---|
operand |
XlaOp |
正規化する N 次元配列 |
scale |
XlaOp |
1 次元配列 |
offset |
XlaOp |
1 次元配列 |
mean |
XlaOp |
1 次元配列 |
variance |
XlaOp |
1 次元配列 |
epsilon |
float |
Epsilon 値 |
feature_index |
int64 |
operand の特徴ディメンションのインデックス |
特徴ディメンションの各特徴(feature_index
は operand
の特徴ディメンションのインデックス)に対して、他のすべてのディメンションの平均と分散を計算し、平均と分散を使用して operand
の各要素を正規化します。feature_index
には、operand
内の特徴ディメンションの有効なインデックスを指定する必要があります。
BatchNormInference
は、バッチごとに mean
と variance
を計算せずに BatchNormTraining
を呼び出す場合と同じです。代わりに、入力 mean
と variance
を推定値として使用します。この op の目的は推論のレイテンシを短縮することであるため、BatchNormInference
という名前が付けられています。
出力は、入力 operand
と同じ形状の N 次元の正規化配列です。
BatchNormTraining
アルゴリズムの詳細な説明については、XlaBuilder::BatchNormTraining
と the original batch normalization paper
をご覧ください。
バッチ ディメンションと空間ディメンション全体で配列を正規化します。
BatchNormTraining(operand, scale, offset, epsilon, feature_index)
引数 | タイプ | セマンティクス |
---|---|---|
operand |
XlaOp |
正規化する n 次元配列(x) |
scale |
XlaOp |
1 次元配列(\(\gamma\)) |
offset |
XlaOp |
1 次元配列(\(\beta\)) |
epsilon |
float |
イプシロン値(\(\epsilon\)) |
feature_index |
int64 |
operand の特徴ディメンションのインデックス |
特徴ディメンションの各特徴(feature_index
は operand
の特徴ディメンションのインデックス)に対して、他のすべてのディメンションの平均と分散を計算し、平均と分散を使用して operand
の各要素を正規化します。feature_index
には、operand
内の特徴ディメンションの有効なインデックスを指定する必要があります。
空間次元のサイズとして w
と h
を持つ m
要素を含む operand
\(x\) の各バッチに対して、アルゴリズムは次のようになります(operand
が 4 次元配列の場合)。
特徴ディメンションの特徴
l
ごとにバッチ平均 \(\mu_l\) を計算します。\(\mu_l=\frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h x_{ijkl}\)バッチ分散を計算します \(\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$
正規化、スケーリング、シフト:\(y_{ijkl}=\frac{\gamma_l(x_{ijkl}-\mu_l)}{\sqrt[2]{\sigma^2_l+\epsilon} }+\beta_l\)
イプシロン値(通常は小さな数)は、ゼロ除算エラーを回避するために加算されます。
出力型は 3 つの XlaOp
のタプルです。
出力 | タイプ | セマンティクス |
---|---|---|
output
|
XlaOp
|
入力 operand (y)と同じ形状の n 次元配列 |
batch_mean |
XlaOp |
1 次元配列(\(\mu\)) |
batch_var |
XlaOp |
1 次元配列(\(\sigma^2\)) |
batch_mean
と batch_var
は、上記の式を使用してバッチ ディメンションと空間ディメンション全体で計算されたモーメントです。
BitcastConvertType
XlaBuilder::BitcastConvertType
もご覧ください。
TensorFlow の tf.bitcast
と同様に、データシェイプからターゲット シェイプへの要素単位のビットキャスト演算を実行します。入力サイズと出力サイズが一致している必要があります。たとえば、s32
要素はビットキャスト ルーチンによって f32
要素になり、1 つの s32
要素は 4 つの s8
要素になります。ビットキャストは低レベルのキャストとして実装されているため、マシンの浮動小数点表現が異なると結果は異なります。
BitcastConvertType(operand, new_element_type)
引数 | タイプ | セマンティクス |
---|---|---|
operand |
XlaOp |
ディメンション D を持つ T 型の配列 |
new_element_type |
PrimitiveType |
U 型 |
オペランドのディメンションとターゲットシェイプのディメンションは、変換前後のプリミティブ サイズの比率に応じて変化する最終ディメンションを除き、一致している必要があります。
コピー元とコピー先の要素の型をタプルにすることはできません。
幅の異なるプリミティブ型へのビットキャスト変換
BitcastConvert
HLO 命令は、出力要素型 T'
のサイズが入力要素 T
のサイズと等しくない場合をサポートします。演算全体は概念的にはビットキャストであり、基になるバイトは変更されないため、出力要素の形状を変更する必要があります。B = sizeof(T), B' =
sizeof(T')
の場合、次の 2 つが考えられます。
まず、B > B'
の場合、出力シェイプはサイズ B/B'
の新しいマイナーなディメンションを取得します。次に例を示します。
f16[10,2]{1,0} %output = f16[10,2]{1,0} bitcast-convert(f32[10]{0} %input)
有効なスカラーのルールは変わりません。
f16[2]{0} %output = f16[2]{0} bitcast-convert(f32[] %input)
あるいは、B' > B
の場合、入力シェイプの最後の論理ディメンションが B'/B
と等しくなる必要があり、このディメンションは変換中に破棄されます。
f32[10]{0} %output = f32[10]{0} bitcast-convert(f16[10,2]{1,0} %input)
なお、異なるビット幅間の変換は要素単位ではありません。
配信
XlaBuilder::Broadcast
もご覧ください。
配列内のデータを複製して、配列にディメンションを追加します。
Broadcast(operand, broadcast_sizes)
引数 | タイプ | セマンティクス |
---|---|---|
operand |
XlaOp |
複製する配列 |
broadcast_sizes |
ArraySlice<int64> |
新しいディメンションのサイズ |
新しいディメンションは左側に挿入されます。つまり、broadcast_sizes
の値が {a0, ..., aN}
で、オペランド形状のディメンションが {b0, ..., bM}
の場合、出力の形状のディメンションは {a0, ..., aN, b0, ..., bM}
になります。
新しいディメンションは、オペランドのコピーにインデックスを付けます。つまり、
output[i0, ..., iN, j0, ..., jM] = operand[j0, ..., jM]
たとえば、operand
が値 2.0f
のスカラー f32
で、broadcast_sizes
が {2, 3}
の場合、結果は形状 f32[2, 3]
の配列となり、結果のすべての値は 2.0f
になります。
BroadcastInDim
XlaBuilder::BroadcastInDim
もご覧ください。
配列内のデータを複製して、配列のサイズとランクを拡張します。
BroadcastInDim(operand, out_dim_size, broadcast_dimensions)
引数 | タイプ | セマンティクス |
---|---|---|
operand |
XlaOp |
複製する配列 |
out_dim_size |
ArraySlice<int64> |
ターゲット シェイプのサイズ |
broadcast_dimensions |
ArraySlice<int64> |
ターゲットのシェイプでオペランドのシェイプの各次元が対応する次元 |
Broadcast と似ていますが、任意の場所にディメンションを追加したり、既存のディメンションをサイズ 1 に拡張したりできます。
operand
は、out_dim_size
で記述される形状にブロードキャストされます。broadcast_dimensions
は、operand
のディメンションをターゲット シェイプのディメンションにマッピングします。つまり、オペランドの i 次元は出力シェイプの broadcast_dimension[i] 次元にマッピングされます。operand
のディメンションは 1 にするか、マッピング先の出力シェイプのディメンションと同じサイズにする必要があります。残りのディメンションにはサイズ 1 のディメンションが入力されます。次に、縮退次元ブロードキャストが、これらの縮退次元に沿ってブロードキャストして、出力シェイプに到達させます。セマンティクスの詳細については、ブロードキャストのページをご覧ください。
電話
XlaBuilder::Call
もご覧ください。
指定された引数で計算を呼び出します。
Call(computation, args...)
引数 | タイプ | セマンティクス |
---|---|---|
computation |
XlaComputation |
任意の型の N 個のパラメータを使用した T_0, T_1, ..., T_{N-1} -> S 型の計算 |
args |
N 個の XlaOp のシーケンス |
任意の型の N 個の引数 |
args
のアリティと型は、computation
のパラメータと一致する必要があります。args
を指定することはできません。
コレスキー
XlaBuilder::Cholesky
もご覧ください。
対称(エルミート)正定行列のバッチのコレスキー分解を計算します。
Cholesky(a, lower)
引数 | タイプ | セマンティクス |
---|---|---|
a |
XlaOp |
複素数または浮動小数点型の rank > 2 の配列。 |
lower |
bool |
a の上または下の三角形を使用するかどうか。 |
lower
が true
の場合、$a = l となるように下三角行列 l
を計算します。
l^T$。lower
が false
の場合、\(a = u^T . u\)となるように上三角行列 u
を計算します。
入力データは、lower
の値に応じて、a
の下/上三角形からのみ読み取られます。もう一方の三角形の値は無視されます。出力データは同じ三角形内に返されます。もう一方の三角形の値は実装で定義されるもので、任意の値を指定できます。
a
のランクが 2 より大きい場合、a
は行列のバッチとして扱われ、マイナー 2 次元を除くすべての次元がバッチ ディメンションです。
a
が対称(エルミート)正定値でない場合、結果は実装で定義されます。
クランプ
XlaBuilder::Clamp
もご覧ください。
オペランドを最小値と最大値の間の範囲内に収めます。
Clamp(min, operand, max)
引数 | タイプ | セマンティクス |
---|---|---|
min |
XlaOp |
T 型の配列 |
operand |
XlaOp |
T 型の配列 |
max |
XlaOp |
T 型の配列 |
オペランドと、最小値と最大値を指定すると、最小値と最大値の間の範囲内であればオペランドを返します。オペランドがこの範囲より下にある場合は最小値を返します。オペランドがこの範囲を超える場合は最大値を返します。つまり、clamp(a, x, b) = min(max(a, x), b)
のようになります。
3 つの配列はすべて同じ形状である必要があります。また、ブロードキャストの制限付き形式として、min
や max
を T
型のスカラーにすることもできます。
スカラー min
と max
の例:
let operand: s32[3] = {-1, 5, 9};
let min: s32 = 0;
let max: s32 = 6;
==>
Clamp(min, operand, max) = s32[3]{0, 5, 6};
閉じる
XlaBuilder::Collapse
オペレーションと tf.reshape
オペレーションもご覧ください。
配列のディメンションを 1 つのディメンションに折りたたみます。
Collapse(operand, dimensions)
引数 | タイプ | セマンティクス |
---|---|---|
operand |
XlaOp |
T 型の配列 |
dimensions |
int64 個のベクトル |
順序付けされた、T の次元の連続するサブセットを取得します。 |
折りたたみは、オペランドのディメンションの特定のサブセットを 1 つのディメンションに置き換えます。入力引数は、T 型の任意の配列と、ディメンション インデックスのコンパイル時定数ベクトルです。ディメンション インデックスは、T のディメンションの連続するサブセットであり、順不同(低から高の次元番号)である必要があります。したがって、{0, 1, 2}、{0, 1}、{1, 2} はすべて有効なディメンション セットですが、{1, 0} または {0, 2} は無効です。これらは、元のディメンション サイズの積と同じ新しいディメンション サイズで、ディメンション シーケンスで置き換えられる場所と同じ位置にある 1 つの新しいディメンションに置き換えられます。dimensions
の最小ディメンション番号は、これらのディメンションを折りたたむループネスト内で最も変化が遅いディメンション(最もメジャー)であり、最大のディメンション番号は変化が最も速いディメンション(最もマイナー)です。より一般的な折りたたみの順序が必要な場合は、tf.reshape
演算子をご覧ください。
たとえば、v が 24 の要素からなる配列であるとします。
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
XlaBuilder::CollectivePermute
もご覧ください。
CollectivePermute は、レプリカ間でデータを送受信する集団オペレーションです。
CollectivePermute(operand, source_target_pairs)
引数 | タイプ | セマンティクス |
---|---|---|
operand |
XlaOp |
n 次元の入力配列 |
source_target_pairs |
<int64, int64> 個のベクトル |
source_replica_id、target_replica_id のペアのリスト。ペアごとに、オペランドがソースレプリカからターゲット レプリカに送信されます。 |
source_target_pair
には次の制限があります。
- 2 つのペアのターゲット レプリカ ID は同じではなく、ソース レプリカ ID も同じであってはなりません。
- レプリカ ID がどのペアのターゲットでもない場合、そのレプリカの出力は入力と同じ形状の 0 から構成されるテンソルになります。
Concatenate
XlaBuilder::ConcatInDim
もご覧ください。
Concatenate は、複数の配列オペランドから配列を作成します。配列は、入力配列オペランドのそれぞれと同じランク(互いに同じランクでなければなりません)であり、指定された順序で引数を含みます。
Concatenate(operands..., dimension)
引数 | タイプ | セマンティクス |
---|---|---|
operands |
N 個の XlaOp のシーケンス |
T 型で次元が [L0, L1, ...] の配列 N 個。N >= 1 が必要です。 |
dimension |
int64 |
operands 間で連結されるディメンションの名前を指定する、[0, N) 区間の値。 |
dimension
を除き、すべてのディメンションを同じにする必要があります。これは、XLA が「不規則な」配列をサポートしていないためです。また、ランク 0 値は連結できないことにも注意してください(連結が行われるディメンションに名前を付けることができないため)。
1 次元の例:
Concat({ {2, 3}, {4, 5}, {6, 7} }, 0)
>>> {2, 3, 4, 5, 6, 7}
2 次元の例:
let a = {
{1, 2},
{3, 4},
{5, 6},
};
let b = {
{7, 8},
};
Concat({a, b}, 0)
>>> {
{1, 2},
{3, 4},
{5, 6},
{7, 8},
}
図:
条件
XlaBuilder::Conditional
もご覧ください。
Conditional(pred, true_operand, true_computation, false_operand,
false_computation)
引数 | タイプ | セマンティクス |
---|---|---|
pred |
XlaOp |
PRED 型のスカラー |
true_operand |
XlaOp |
\(T_0\)型の引数 |
true_computation |
XlaComputation |
\(T_0 \to S\)型の XlaComputation |
false_operand |
XlaOp |
\(T_1\)型の引数 |
false_computation |
XlaComputation |
\(T_1 \to S\)型の XlaComputation |
pred
が true
の場合は true_computation
、pred
が false
の場合は false_computation
を実行し、結果を返します。
true_computation
は、 \(T_0\) 型の引数を 1 つ受け取る必要があり、同じ型である true_operand
で呼び出されます。false_computation
は、 \(T_1\) 型の引数を 1 つ受け取る必要があり、同じ型である false_operand
で呼び出されます。true_computation
と false_computation
の戻り値の型は同じである必要があります。
pred
の値に応じて、true_computation
と false_computation
のいずれか 1 つのみが実行されることに注意してください。
Conditional(branch_index, branch_computations, branch_operands)
引数 | タイプ | セマンティクス |
---|---|---|
branch_index |
XlaOp |
S32 型のスカラー |
branch_computations |
N 個の XlaComputation のシーケンス |
\(T_0 \to S , T_1 \to S , ..., T_{N-1} \to S\)型の XlaComputations |
branch_operands |
N 個の XlaOp のシーケンス |
\(T_0 , T_1 , ..., T_{N-1}\)型の引数 |
branch_computations[branch_index]
を実行し、結果を返します。branch_index
が 0 未満または N 以上の S32
の場合、branch_computations[N-1]
がデフォルトの分岐として実行されます。
各 branch_computations[b]
は、 \(T_b\) 型の引数を 1 つ受け取る必要があり、同じ型である branch_operands[b]
で呼び出されます。各 branch_computations[b]
の戻り値の型は同じである必要があります。
branch_index
の値に応じて、branch_computations
の 1 つのみが実行されます。
コンバージョン(畳み込み)
XlaBuilder::Conv
もご覧ください。
ConvWithGeneralPadding と同じですが、パディングは簡潔に SAME または VALID として指定されます。SAME パディングは、入力(lhs
)をゼロでパディングして、ストライドを考慮しない場合の出力が入力と同じ形状になるようにします。VALID パディングとは、単にパディングがないことを意味します。
ConvWithGeneralPadding(畳み込み)
XlaBuilder::ConvWithGeneralPadding
もご覧ください。
ニューラル ネットワークで使用される種類の畳み込みを計算します。ここで、畳み込みは n 次元のベース領域上を移動する n 次元のウィンドウと考えることができ、ウィンドウの考えられる位置ごとに計算が実行されます。
引数 | タイプ | セマンティクス |
---|---|---|
lhs |
XlaOp |
入力配列のランク n+2 |
rhs |
XlaOp |
カーネルの重みのランク n+2 配列 |
window_strides |
ArraySlice<int64> |
カーネル ストライドの n 次元配列 |
padding |
ArraySlice< pair<int64,int64>> |
パディング(低、高)の n 配列 |
lhs_dilation |
ArraySlice<int64> |
n-d lhs 拡張係数配列 |
rhs_dilation |
ArraySlice<int64> |
n-d rhs 拡張係数配列 |
feature_group_count |
int64 | 特徴グループの数 |
batch_group_count |
int64 | バッチグループの数 |
n を空間次元の数とします。lhs
引数は、基本領域を表すランク n+2 の配列です。これは入力と呼ばれますが
もちろん Rhs は入力でもありますニューラル ネットワークでは、これらは入力の活性化です。n+2 次元は次の順序です。
batch
: このディメンションの各座標は、畳み込みを行う独立した入力を表します。z/depth/features
: ベース領域内の各(y、x)位置には、この次元に入るベクトルが関連付けられています。spatial_dims
: ウィンドウが移動する基本領域を定義するn
空間寸法を記述します。
rhs
引数は、畳み込みフィルタ/カーネル/ウィンドウを記述するランク n+2 配列です。ディメンションは次の順序で表示されます。
output-z
: 出力のz
ディメンション。input-z
: このディメンションのサイズにfeature_group_count
を掛けた値は、z
ディメンションのサイズに LHS で等しい必要があります。spatial_dims
: ベース領域内を移動する n-d ウィンドウを定義するn
空間寸法を記述します。
window_strides
引数は、空間次元での畳み込みウィンドウのストライドを指定します。たとえば、最初の空間次元のストライドが 3 の場合、ウィンドウは最初の空間インデックスが 3 で割り切れる座標にのみ配置できます。
padding
引数は、ベース領域に適用するゼロパディングの量を指定します。パディングの量は負の値にできます。負のパディングの絶対値は、畳み込み処理を行う前に、指定されたディメンションから削除する要素の数を示します。padding[0]
は y
のディメンションのパディングを指定し、padding[1]
は x
のディメンションのパディングを指定します。各ペアでは、最初の要素として低いパディング、2 番目の要素として高いパディングがあります。低いパディングは低いインデックスの方向に適用され、高いパディングは高いインデックスの方向に適用されます。たとえば、padding[1]
が (2,3)
の場合、2 番目の空間次元では左側に 2 個のゼロ、右側に 3 個のゼロのパディングがあります。パディングを使用することは、畳み込み処理を行う前に同じゼロ値を入力(lhs
)に挿入するのと同じです。
lhs_dilation
引数と rhs_dilation
引数は、各空間次元で lhs と rhs に適用する拡張係数を指定します。空間次元の拡張係数を d とすると、その次元の各エントリの間に d-1 ホールが暗黙的に配置され、配列のサイズが大きくなります。穴には no-op 値が入力されます。畳み込み演算では 0 になります。
右領域の拡張は、atrous 畳み込みとも呼ばれます。詳しくは、tf.nn.atrous_conv2d
をご覧ください。lhs の拡張は転置畳み込みとも呼ばれます。詳しくは、tf.nn.conv2d_transpose
をご覧ください。
グループ化された畳み込みには、feature_group_count
引数(デフォルト値 1)を使用できます。feature_group_count
は、入力と出力の両方の特徴ディメンションの除数である必要があります。feature_group_count
が 1 より大きい場合、概念的には、入出力特徴ディメンションと rhs
出力特徴ディメンションが、多くの feature_group_count
グループに均等に分割され、各グループは連続する特徴のサブシーケンスで構成されます。rhs
の入力特徴ディメンションは、lhs
入力特徴ディメンションを feature_group_count
で割ったものに等しくする必要があります(入力特徴のグループのサイズにすでに相当するため)。i 番目のグループは、多くの個別の畳み込みの feature_group_count
を計算するために一緒に使用されます。これらの畳み込みの結果は、出力特徴ディメンションで連結されます。
深度方向の畳み込みでは、feature_group_count
引数が入力特徴次元に設定され、フィルタの形状が [filter_height, filter_width, in_channels, channel_multiplier]
から [filter_height, filter_width, 1, in_channels * channel_multiplier]
に再形成されます。詳しくは、tf.nn.depthwise_conv2d
をご覧ください。
batch_group_count
(デフォルト値 1)引数は、バックプロパゲーション中にグループ化フィルタに使用できます。batch_group_count
は、lhs
(入力)バッチ ディメンションのサイズの除数である必要があります。batch_group_count
が 1 より大きい場合、出力バッチ ディメンションはサイズ input batch
/ batch_group_count
にする必要があります。batch_group_count
は、出力特徴サイズの除数にする必要があります。
出力シェイプは、次の順序で表示されます。
batch
: このディメンションのサイズにbatch_group_count
を掛けた値は、batch
ディメンションのサイズと lh 単位で一致する必要があります。z
: カーネル(rhs
)上のoutput-z
と同じサイズ。spatial_dims
: 畳み込みウィンドウの有効な配置ごとに 1 つの値。
上の図は、batch_group_count
フィールドの仕組みを示しています。実質的には、各 lhs バッチを batch_group_count
グループにスライスし、出力特徴についても同じことを行います。次に、これらのグループごとに、ペアワイズ畳み込みを行い、出力特徴ディメンションに沿って出力を連結します。他のすべてのディメンション(特徴量と空間)の運用セマンティクスは変わりません。
畳み込みウィンドウの有効な配置は、ストライドと、パディング後のベース領域のサイズによって決まります。
畳み込みが何をするのかを説明するには、2 次元畳み込みについて考え、出力に固定された batch
、z
、y
、x
座標を選択します。この場合、(y,x)
はベース領域内のウィンドウの角の位置です(空間次元の解釈方法によっては左上の角など)。これで、ベース領域から取得した 2D ウィンドウができあがります。各 2D のポイントが 1D ベクトルに関連付けられるため、3D のボックスができあがります。畳み込みカーネルで出力座標 z
を固定したため、3D ボックスも出力されます。2 つのボックスの寸法は同じであるため、2 つのボックス間で要素ごとの積の合計を求めることができます(ドット積と同様)。これが出力値です。
なお、output-z
が例の場合、ウィンドウの各位置から、出力の z
ディメンションへの出力に 5 つの値が生成されます。これらの値は、畳み込みカーネルのどの部分を使用するかによって異なります。各 output-z
座標に使用される値の 3D ボックスが個別に存在します。それぞれに対して異なるフィルタを使用した 5 つの異なる畳み込みと考えることができます。
パディングとストライディングを使用した 2 次元畳み込みの擬似コードは次のとおりです。
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
XlaBuilder::ConvertElementType
もご覧ください。
C++ の要素単位の static_cast
と同様に、データシェイプからターゲットシェイプへの要素単位の変換オペレーションを実行します。ディメンションが一致する必要があり、変換は要素ごとに行われます。たとえば、s32
要素は、s32
から f32
への変換ルーティンによって f32
要素になります。
ConvertElementType(operand, new_element_type)
引数 | タイプ | セマンティクス |
---|---|---|
operand |
XlaOp |
ディメンション D を持つ T 型の配列 |
new_element_type |
PrimitiveType |
U 型 |
オペランドの寸法とターゲットの形状は一致している必要があります。ソース要素タイプとデスティネーション要素タイプをタプルにすることはできません。
T=s32
から U=f32
への変換では、int から float への変換ルーチン(たとえば、round-nearest-even など)が実行されます。
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
総和計算を使用して AllReduce
を実行します。
CustomCall
XlaBuilder::CustomCall
もご覧ください。
計算内でユーザー指定の関数を呼び出します。
CustomCall(target_name, args..., shape)
引数 | タイプ | セマンティクス |
---|---|---|
target_name |
string |
関数の名前。このシンボル名をターゲットとする呼び出し命令が出されます。 |
args |
N 個の XlaOp のシーケンス |
任意の型の N 個の引数。関数に渡されます。 |
shape |
Shape |
関数の出力形状 |
引数のアリティや型に関係なく、関数のシグネチャは同じです。
extern "C" void target_name(void* out, void** in);
たとえば、次のように CustomCall を使用するとします。
let x = f32[2] {1,2};
let y = f32[2x3] { {10, 20, 30}, {40, 50, 60} };
CustomCall("myfunc", {x, y}, f32[3x3])
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];
// ...
}
ユーザー指定関数に副作用があってはなりません。また、その実行はべき等でなければなりません。
Dot
XlaBuilder::Dot
もご覧ください。
Dot(lhs, rhs)
引数 | タイプ | セマンティクス |
---|---|---|
lhs |
XlaOp |
T 型の配列 |
rhs |
XlaOp |
T 型の配列 |
この演算の正確なセマンティクスは、オペランドのランクによって異なります。
入力 | 出力 | セマンティクス |
---|---|---|
ベクトル [n] dot ベクトル [n] |
スカラー | ベクトル ドット積 |
行列 [m x k] dot ベクトル [k] |
ベクトル [m] | 行列ベクトル乗算 |
行列 [m x k] dot 行列 [k x n] |
行列 [m x n] | 行列行列の乗算 |
この演算は、lhs
の 2 つ目の次元(ランク 1 の場合は最初の次元)と rhs
の 1 つ目の次元の積の合計を実行します。これらは「契約」ディメンションです。lhs
と rhs
の契約サイズは同じサイズにする必要があります。実際には、ベクトル間のドット積、ベクトル/行列の乗算、行列/行列の乗算に使用できます。
DotGeneral
XlaBuilder::DotGeneral
もご覧ください。
DotGeneral(lhs, rhs, dimension_numbers)
引数 | タイプ | セマンティクス |
---|---|---|
lhs |
XlaOp |
T 型の配列 |
rhs |
XlaOp |
T 型の配列 |
dimension_numbers |
DotDimensionNumbers |
分割とバッチのディメンション番号 |
Dot と似ていますが、lhs
と rhs
の両方に縮小ディメンションとバッチ ディメンション番号を指定できます。
DotDimensionNumbers フィールド | タイプ | セマンティクス |
---|---|---|
lhs_contracting_dimensions
|
繰り返し int64 | lhs 契約ディメンション番号 |
rhs_contracting_dimensions
|
繰り返し int64 | rhs 契約ディメンション番号 |
lhs_batch_dimensions
|
繰り返し int64 | lhs 個のバッチ ディメンション番号 |
rhs_batch_dimensions
|
繰り返し int64 | rhs 個のバッチ ディメンション番号 |
DotGeneral は、dimension_numbers
で指定された分割ディメンションに対する積の合計を実行します。
lhs
と rhs
の関連付けられた縮小ディメンション番号は同じにする必要はありませんが、同じディメンション サイズにする必要があります。
ディメンション番号を縮小する例:
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} }
lhs
と rhs
の関連するバッチ ディメンション番号は、同じディメンション サイズである必要があります。
バッチ ディメンション番号の例(バッチサイズ 2、2×2 行列):
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} } }
入力 | 出力 | セマンティクス |
---|---|---|
[b0, m, k] dot [b0, k, n] |
[b0, m, n] | バッチ matmul |
[b0, b1, m, k] dot [b0, b1, k, n] |
[b0, b1, m, n] | バッチ matmul |
したがって、結果のディメンション番号はバッチ ディメンションで始まり、lhs
非契約/非バッチ ディメンション、最後に rhs
非契約/非バッチ ディメンションと続きます。
DynamicSlice
XlaBuilder::DynamicSlice
もご覧ください。
DynamicSlice は、動的 start_indices
の入力配列からサブ配列を抽出します。各ディメンションのスライスのサイズは size_indices
で渡され、各ディメンションの排他的スライス間隔の終了点([start, start + size])を指定します。start_indices
の形状はランク == 1、ディメンション サイズは operand
のランクにする必要があります。
DynamicSlice(operand, start_indices, size_indices)
引数 | タイプ | セマンティクス |
---|---|---|
operand |
XlaOp |
T 型の N 次元配列 |
start_indices |
N 個の XlaOp のシーケンス |
各次元のスライスの開始インデックスを含む N 個のスカラー整数のリスト。0 以上の値を指定してください。 |
size_indices |
ArraySlice<int64> |
各ディメンションのスライスサイズを含む N 個の整数のリスト。モジュロ ディメンション サイズがラップされないように、各値は厳密に 0 より大きく、start + size はディメンションのサイズ以下にする必要があります。 |
有効なスライス インデックスは、スライスを実行する前に、[1, N)
のインデックス i
ごとに次の変換を適用して計算されます。
start_indices[i] = clamp(start_indices[i], 0, operand.dimension_size[i] - size_indices[i])
これにより、抽出されたスライスは常にオペランド配列の範囲内になります。変換が適用される前にスライスが範囲内にある場合、変換の効果はありません。
1 次元の例:
let a = {0.0, 1.0, 2.0, 3.0, 4.0}
let s = {2}
DynamicSlice(a, s, {2}) produces:
{2.0, 3.0}
2 次元の例:
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
XlaBuilder::DynamicUpdateSlice
もご覧ください。
DynamicUpdateSlice は、入力配列 operand
の値である結果を生成します。スライス update
は start_indices
で上書きされます。update
の形状によって、更新される結果のサブ配列の形状が決まります。
start_indices
の形状はランク == 1、ディメンション サイズは operand
のランクにする必要があります。
DynamicUpdateSlice(operand, update, start_indices)
引数 | タイプ | セマンティクス |
---|---|---|
operand |
XlaOp |
T 型の N 次元配列 |
update |
XlaOp |
スライスの更新を含む T 型の N 次元配列。範囲外の更新インデックスが生成されないように、更新シェイプの各ディメンションは厳密に 0 より大きく、start + update は各ディメンションのオペランド サイズ以下にする必要があります。 |
start_indices |
N 個の XlaOp のシーケンス |
各次元のスライスの開始インデックスを含む N 個のスカラー整数のリスト。0 以上の値を指定してください。 |
有効なスライス インデックスは、スライスを実行する前に、[1, N)
のインデックス i
ごとに次の変換を適用して計算されます。
start_indices[i] = clamp(start_indices[i], 0, operand.dimension_size[i] - update.dimension_size[i])
これにより、更新されたスライスは常にオペランド配列の範囲内になります。変換が適用される前にスライスが範囲内にある場合、変換の効果はありません。
1 次元の例:
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}
2 次元の例:
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} }
要素ごとの 2 項算術演算
XlaBuilder::Add
もご覧ください。
要素単位のバイナリ算術演算のセットがサポートされています。
Op(lhs, rhs)
ここで、Op
は Add
(加算)、Sub
(減算)、Mul
(乗算)、Div
(除算)、Rem
(剰余)、Max
(最大)、Min
(最小)、LogicalAnd
(論理 AND)、LogicalOr
(論理 OR)のいずれかです。
引数 | タイプ | セマンティクス |
---|---|---|
lhs |
XlaOp |
左側のオペランド: T 型の配列 |
rhs |
XlaOp |
右側のオペランド: T 型の配列 |
引数の形状は類似しているか、互換性がある必要があります。シェイプの互換性の意味については、ブロードキャストのドキュメントをご覧ください。演算の結果の形状は、2 つの入力配列をブロードキャストした結果です。このバリアントでは、オペランドの 1 つがスカラーでない限り、異なるランクの配列間の演算はサポートされません。
Op
が Rem
の場合、結果の符号は被除数から取得され、結果の絶対値は常に除数の絶対値よりも小さくなります。
整数除算オーバーフロー(ゼロによる符号付き/符号なし除算/剰余、または -1
を使用した INT_SMIN
の符号付き除算/剰余)は、実装で定義された値を生成します。
次のオペレーションには、異なるランクのブロードキャストをサポートする別のバリアントがあります。
Op(lhs, rhs, broadcast_dimensions)
ここで、Op
は上記と同じです。この演算のバリアントは、異なるランクの配列間の算術演算(ベクトルへの行列の追加など)に使用します。
追加の broadcast_dimensions
オペランドは、低ランクオペランドのランクを高ランクオペランドのランクに拡張するために使用される整数のスライスです。broadcast_dimensions
は、下位シェイプのディメンションを上位シェイプのディメンションにマッピングします。展開されたシェイプのマッピングされていないサイズは、サイズ 1 のディメンションで埋められます。次に、縮退次元ブロードキャストは、これらの縮退次元に沿って形状をブロードキャストして、両方のオペランドの形状を均等にします。セマンティクスの詳細については、ブロードキャストのページをご覧ください。
要素単位の比較演算
XlaBuilder::Eq
もご覧ください。
標準的な要素単位のバイナリ比較演算のセットがサポートされています。浮動小数点型を比較する場合は、標準の IEEE 754 浮動小数点比較セマンティクスが適用されます。
Op(lhs, rhs)
ここで、Op
は Eq
(等しい)、Ne
(等しくない)、Ge
(次より大きい)、Gt
(次より大きい)、Le
(次より小さい)、Lt
(次より小さい)のいずれかです。別の演算子セット EqTotalOrder、NeTotalOrder、GeTotalOrder、GtTotalOrder、LeTotalOrder、LtTotalOrder は同じ機能を提供しますが、-NaN < -Inf < -Finite < -0 < +0 < +InaN < +Inf < +Finite を適用することで、浮動小数点数に対する合計順序もサポートします。
引数 | タイプ | セマンティクス |
---|---|---|
lhs |
XlaOp |
左側のオペランド: T 型の配列 |
rhs |
XlaOp |
右側のオペランド: T 型の配列 |
引数の形状は類似しているか、互換性がある必要があります。シェイプの互換性の意味については、ブロードキャストのドキュメントをご覧ください。演算の結果の形状は、要素型 PRED
で 2 つの入力配列をブロードキャストした結果です。このバリアントでは、オペランドの 1 つがスカラーでない限り、異なるランクの配列間の演算はサポートされません。
次のオペレーションには、異なるランクのブロードキャストをサポートする別のバリアントがあります。
Op(lhs, rhs, broadcast_dimensions)
ここで、Op
は上記と同じです。この演算のバリアントは、異なるランクの配列間の比較演算(ベクトルへの行列の追加など)に使用します。
追加の broadcast_dimensions
オペランドは、オペランドのブロードキャストに使用するディメンションを指定する整数のスライスです。セマンティクスの詳細については、ブロードキャスト ページをご覧ください。
要素単位の単項関数
XlaBuilder は、次の要素単位の単項関数をサポートしています。
Abs(operand)
要素単位の絶対値 x -> |x|
。
Ceil(operand)
要素単位の ceil x -> ⌈x⌉
。
Cos(operand)
要素単位のコサイン x -> cos(x)
。
Exp(operand)
要素ごとの自然指数 x -> e^x
。
Floor(operand)
要素単位の最小価格 x -> ⌊x⌋
。
Imag(operand)
複雑な(または実数)図形の要素単位の虚数部。x -> imag(x)
。オペランドが浮動小数点型の場合は 0 を返します。
IsFinite(operand)
operand
の各要素が有限かどうか(つまり、正の無限大でも負の無限大でもなく、NaN
でもない)かどうかをテストします。入力と同じ形状の PRED
値の配列を返します。各要素は、対応する入力要素が有限である場合にのみ true
になります。
Log(operand)
要素単位の自然対数 x -> ln(x)
。
LogicalNot(operand)
x -> !(x)
ではなく、要素ごとの論理値です。
Logistic(operand)
要素単位のロジスティック関数の計算 x ->
logistic(x)
。
PopulationCount(operand)
operand
の各要素に設定されているビット数を計算します。
Neg(operand)
要素単位の否定の x -> -x
。
Real(operand)
複雑な(または実際の)シェイプの要素単位の実数部。
x -> real(x)
。オペランドが浮動小数点型の場合、同じ値を返します。
Rsqrt(operand)
x -> 1.0 / sqrt(x)
平方根演算の要素単位の逆数。
Sign(operand)
要素ごとの符号演算 x -> sgn(x)
。ここで
\[\text{sgn}(x) = \begin{cases} -1 & x < 0\\ -0 & x = -0\\ NaN & x = NaN\\ +0 & x = +0\\ 1 & x > 0 \end{cases}\]
operand
の要素の型の比較演算子を使用します。
Sqrt(operand)
要素ごとの平方根演算 x -> sqrt(x)
です。
Cbrt(operand)
要素単位の立方根演算 x -> cbrt(x)
。
Tanh(operand)
要素ごとの双曲線正接の x -> tanh(x)
。
Round(operand)
要素ごとの丸め、ゼロから離れる同値。
RoundNearestEven(operand)
要素単位で丸め、最も近い偶数に関連付ける。
引数 | タイプ | セマンティクス |
---|---|---|
operand |
XlaOp |
関数のオペランド |
この関数は operand
配列内の各要素に適用され、同じ形状の配列になります。operand
はスカラー(ランク 0)にできます。
FFT
XLA FFT 演算は、実数と複雑な入出力に対して順方向と逆方向のフーリエ変換を実装します。最大 3 軸の多次元 FFT がサポートされています。
XlaBuilder::Fft
もご覧ください。
引数 | タイプ | セマンティクス |
---|---|---|
operand |
XlaOp |
フーリエ変換する配列。 |
fft_type |
FftType |
下の表をご覧ください。 |
fft_length |
ArraySlice<int64> |
変換される軸の時間ドメインの長さ。これは特に、RFFT(fft_length=[16]) が RFFT(fft_length=[17]) と同じ出力形状を持つため、IRFFT で最も内側の軸のサイズを適正化するために必要です。 |
FftType |
セマンティクス |
---|---|
FFT |
複雑から複雑な FFT を転送する。形状は変わりません。 |
IFFT |
複素数から複素数への逆 FFT。形状は変わりません。 |
RFFT |
実数から複雑な FFT を転送する。fft_length[-1] がゼロ以外の値の場合、最も内側の軸の形状は fft_length[-1] // 2 + 1 に縮小され、ナイキスト周波数以外の変換された信号の逆共役部分は除外されます。 |
IRFFT |
実数から複素数への逆 FFT(複素数を取り、実数を返す)。fft_length[-1] がゼロ以外の値の場合、最も内側の軸の形状は fft_length[-1] に展開され、変換された信号のナイキスト周波数を超える部分は 1 のエントリの逆共役から fft_length[-1] // 2 + 1 のエントリから推測されます。 |
多次元 FFT
複数の fft_length
を指定した場合、最も内側の軸のそれぞれに FFT 演算のカスケードを適用するのと同じ結果になります。実数 - 複素数 - 実数の場合、最も内側の軸の変換が(事実上)最初に(RFFT、IRFFT の場合は最後に)実行されます。そのため、最も内側の軸でサイズが変更されます。他の軸の変換は複雑から複雑になります。
実装の詳細
CPU FFT は Eigen の TensorFFT を基盤としています。GPU FFT は cuFFT を使用します。
収集
XLA 収集演算は、入力配列の複数のスライス(異なる可能性のあるランタイム オフセットで各スライス)をつなぎ合わせます。
一般的なセマンティクス
XlaBuilder::Gather
もご覧ください。より直感的な説明については、以下の「非公式な説明」をご覧ください。
gather(operand, start_indices, offset_dims, collapsed_slice_dims, slice_sizes, start_index_map)
引数 | タイプ | セマンティクス |
---|---|---|
operand |
XlaOp |
収集元の配列。 |
start_indices |
XlaOp |
収集したスライスの開始インデックスを含む配列。 |
index_vector_dim |
int64 |
開始インデックスを「含む」start_indices のディメンション。詳しくは以下をご覧ください。 |
offset_dims |
ArraySlice<int64> |
オペランドからスライス化された配列にオフセットされる出力シェイプのディメンションのセット。 |
slice_sizes |
ArraySlice<int64> |
slice_sizes[i] は、ディメンション i のスライスの境界です。 |
collapsed_slice_dims |
ArraySlice<int64> |
折りたたまれた各スライスのディメンションのセット。これらのディメンションのサイズは 1 である必要があります。 |
start_index_map |
ArraySlice<int64> |
start_indices のインデックスを、オペランドの有効なインデックスにマッピングする方法を記述するマップ。 |
indices_are_sorted |
bool |
インデックスが呼び出し元によって並べ替えられることが保証されているかどうか。 |
便宜上、offset_dims
内にない出力配列のディメンションには batch_dims
というラベルを付けます。
出力はランク batch_dims.size
+ offset_dims.size
の配列です。
operand.rank
は、offset_dims.size
と collapsed_slice_dims.size
の合計と等しくなければなりません。また、slice_sizes.size
は operand.rank
と等しくする必要があります。
index_vector_dim
が start_indices.rank
と等しい場合、暗黙的に start_indices
に末尾の 1
ディメンションがあると見なします(つまり、start_indices
の形状が [6,7]
で、index_vector_dim
が 2
の場合、start_indices
の形状は暗黙的に [6,7,1]
であると見なされます)。
ディメンション i
に沿った出力配列の境界は次のように計算されます。
i
がbatch_dims
に存在する(一部のk
でbatch_dims[k]
に等しい)場合、対応するディメンションの境界がstart_indices.shape
から選択され、index_vector_dim
はスキップされます(k
<index_vector_dim
の場合はstart_indices.shape.dims
[k
]、それ以外の場合はstart_indices.shape.dims
[k
+1
] が選択されます)。i
がoffset_dims
に存在する(一部のk
のoffset_dims
[k
] に等しい)場合、collapsed_slice_dims
を考慮したうえで、slice_sizes
から対応する境界を選択します(つまり、adjusted_slice_sizes
がslice_sizes
で、インデックスcollapsed_slice_dims
の境界を削除したadjusted_slice_sizes
[k
] を選択します)。
特定の出力インデックス Out
に対応するオペランド インデックス In
は、形式的に次のように計算されます。
batch_dims
内のk
に対してG
= {Out
[k
] } とします。G
を使用して、S
[i
] =start_indices
[Combine(G
,i
)] のようにベクトルS
をスライスします。ここで、Combine(A, b) は位置index_vector_dim
で b を A に挿入します。これは、G
が空であっても明確に定義されます。G
が空の場合、S
=start_indices
となります。start_index_map
を使用してS
を分散し、S
を使用してoperand
に開始インデックスS
in
を作成します。より正確に言うと、S
k
<start_index_map.size
の場合、in
[start_index_map
[k
]] =S
[k
]。S
それ以外はin
[_
] =0
。
collapsed_slice_dims
のセットに従ってOut
のオフセット ディメンションでインデックスを分散させることで、operand
へのインデックスO
in
を作成します。より正確に言うと、O
k
<offset_dims.size
の場合、in
[remapped_offset_dims
(k
)] =Out
[offset_dims
[k
]](remapped_offset_dims
は以下で定義)。O
それ以外はin
[_
] =0
。
In
はO
in
+S
in
です。「+」は要素ごとの加算です。
remapped_offset_dims
は、ドメイン [0
, offset_dims.size
)、範囲 [0
, operand.rank
) \ collapsed_slice_dims
を持つ単調関数です。たとえばoffset_dims.size
は 4
、operand.rank
は 6
、collapsed_slice_dims
は {0
、2
}、remapped_offset_dims
は {0
→1
、1
→3
、2
→4
、3
→5
} です。
indices_are_sorted
が true に設定されている場合、XLA は start_indices
がユーザーによって(start_index_map
の昇順で)並べ替えられていると想定できます。そうでない場合、セマンティクスは実装が定義されています。
非公式の説明と例
非公式に、出力配列のすべてのインデックス Out
は、オペランド配列の要素 E
に対応しています。これは次のように計算されます。
Out
のバッチ ディメンションを使用して、start_indices
から開始インデックスを検索します。start_index_map
を使用して、開始インデックス(このサイズは operand.rank より小さい場合があります)をoperand
の「完全な」開始インデックスにマッピングします。完全な開始インデックスを使用して、サイズ
slice_sizes
のスライスを動的にスライスします。collapsed_slice_dims
ディメンションを折りたたんでスライスの形状を変更します。折りたたみスライスのサイズはすべて、境界が 1 である必要があるため、この形状変更は常に有効です。Out
のオフセット ディメンションを使用してこのスライスにインデックスを付け、出力インデックスOut
に対応する入力要素E
を取得します。
以降のすべての例で、index_vector_dim
は start_indices.rank
- 1
に設定されています。index_vector_dim
の値に興味深い値を指定すると、演算自体は変わりませんが、視覚的表現が複雑になります。
上記すべてがどのように組み合わされているかを理解するために、[16,11]
配列からシェイプ [8,6]
の 5 つのスライスを収集する例を見てみましょう。[16,11]
配列内のスライスの位置は、形状 S64[2]
のインデックス ベクトルとして表すことができるため、5 つの位置のセットは S64[5,2]
配列として表すことができます。
収集オペレーションの動作は、出力シェイプのインデックスである [G
、O
0
、O
1
] を取り、次のように入力配列の要素にマッピングするインデックス変換として表すことができます。
まず、G
を使用して collect インデックス配列から(X
,Y
)ベクトルを選択します。出力配列のインデックス [G
、O
0
、O
1
] にある要素は、インデックス [X
+O
0
,Y
+O
1
] にある入力配列の要素です。
slice_sizes
は [8,6]
です。これにより、O0
と O1
の範囲が決定され、スライスの境界が決定されます。
この収集オペレーションは、G
をバッチ ディメンションとするバッチ動的スライスとして機能します。
ギャザー インデックスは多次元の場合があります。たとえば、上記の例のより一般的なバージョンでシェイプ [4,5,2]
の「gather indexs」配列を使用すると、次のようなインデックスが変換されます。
ここでも、バッチ動的スライス G
0
と G
1
をバッチ ディメンションとして機能します。スライスのサイズは [8,6]
のままです。
XLA での収集オペレーションでは、上記で説明した非公式のセマンティクスが次のように一般化されます。
出力シェイプのどのディメンションがオフセット ディメンションかを構成できます(最後の例では、
O
0
、O
1
を含むディメンション)。出力バッチ ディメンション(最後の例のG
0
、G
1
を含むディメンション)は、オフセット ディメンションではない出力ディメンションとして定義されます。出力シェイプに明示的に存在する出力オフセット次元の数は、入力ランクよりも小さくなる場合があります。このような「欠落」しているディメンションは
collapsed_slice_dims
として明示的にリストされ、スライスサイズを1
にする必要があります。スライスサイズが1
であるため、有効なインデックスは0
のみであり、省略してもあいまいさは発生しません。「Collect Indices」配列(最後の例では
X
、Y
)から抽出されたスライスの要素が入力配列ランクよりも少ない場合があり、明示的なマッピングにより、入力と同じランクになるようにインデックスを拡張する方法が指示されます。
最後の例として、(2)と(3)を使用して tf.gather_nd
を実装します。
G
0
と G
1
は、開始インデックスに X
という 1 つの要素しかないことを除き、通常と同様に collect インデックス配列から開始インデックスをスライスするために使用されます。同様に、値が O
0
の出力オフセット インデックスは 1 つだけです。ただし、入力配列のインデックスとして使用される前に、これらは「集約インデックス マッピング」(正式な説明の start_index_map
)と「オフセット マッピング」(正式な記述の remapped_offset_dims
)に従ってそれぞれ [X
、0
] と [0
、O
0
] に展開され、[X
、O
0
{2] の [G
0
1
1] に拡張されます。つまり、出力インデックス [0
G
O
] は [0
1] に拡張されます。つまり、[0
G
O
] はインデックス [0
1] にマップされます。G
1
GatherIndices
tf.gather_nd
このケースの slice_sizes
は [1,11]
です。これは直感的に、collect インデックス配列内のすべてのインデックス X
が行全体を選択し、その結果はこれらすべての行を連結したものであることを意味します。
GetDimensionSize
XlaBuilder::GetDimensionSize
もご覧ください。
オペランドの指定された次元のサイズを返します。オペランドは配列型でなければなりません。
GetDimensionSize(operand, dimension)
引数 | タイプ | セマンティクス |
---|---|---|
operand |
XlaOp |
n 次元の入力配列 |
dimension |
int64 |
ディメンションを指定する、[0, n) 区間の値 |
SetDimensionSize
XlaBuilder::SetDimensionSize
もご覧ください。
XlaOp の特定のディメンションの動的サイズを設定します。オペランドは配列型でなければなりません。
SetDimensionSize(operand, size, dimension)
引数 | タイプ | セマンティクス |
---|---|---|
operand |
XlaOp |
n 次元の入力配列。 |
size |
XlaOp |
ランタイムの動的サイズを表す int32。 |
dimension |
int64 |
ディメンションを指定する、[0, n) 区間の値。 |
コンパイラによって追跡される動的ディメンションを使用して、オペランドを結果として渡します。
パディングされた値は、ダウンストリームのリダクション オペレーションで無視されます。
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
XlaBuilder::GetTupleElement
もご覧ください。
コンパイル時の定数値を持つタプルにインデックスを付けます。
値はコンパイル時の定数にする必要があります。これにより、形状推論が結果の値の型を決定できるようになります。
これは C++ の std::get<int N>(t)
に似ています。概念的には次のようになります。
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.
tf.tuple
もご覧ください。
インフィード
XlaBuilder::Infeed
もご覧ください。
Infeed(shape)
引数 | タイプ | セマンティクス |
---|---|---|
shape |
Shape |
インフィード インターフェースから読み取られたデータの形状。シェイプのレイアウト フィールドは、デバイスに送信されるデータのレイアウトと一致するように設定する必要があります。そうでない場合、動作は未定義です。 |
デバイスの暗黙的なインフィード ストリーミング インターフェースから単一のデータアイテムを読み取り、データを指定されたシェイプとそのレイアウトとして解釈して、データの XlaOp
を返します。1 回の計算で複数のインフィード オペレーションを使用できますが、インフィード オペレーションの順序は合計する必要があります。たとえば、次のコードでは 2 つのインフィードは、while ループ間に依存関係があるため、合計順序になります。
result1 = while (condition, init = init_value) {
Infeed(shape)
}
result2 = while (condition, init = result1) {
Infeed(shape)
}
ネストされたタプルの形状はサポートされていません。空のタプルシェイプの場合、インフィード オペレーションは実質的に何も行われず、デバイスのインフィードからデータを読み取らずに実行されます。
ロッタ
XlaBuilder::Iota
もご覧ください。
Iota(shape, iota_dimension)
大規模なホスト転送の可能性ではなく、デバイス上で定数リテラルをビルドします。指定された形状を持ち、ゼロから始まり、指定されたディメンションに沿って 1 ずつ増分する値を保持する配列を作成します。浮動小数点型の場合、生成される配列は ConvertElementType(Iota(...))
と同等です。ここで、Iota
は整数型で、変換は浮動小数点型です。
引数 | タイプ | セマンティクス |
---|---|---|
shape |
Shape |
Iota() によって作成された配列の形状 |
iota_dimension |
int64 |
インクリメントするディメンション。 |
たとえば、Iota(s32[4, 8], 0)
は以下を返します。
[[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 ]]
返品可能(返品手数料: 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 ]]
マッピング
XlaBuilder::Map
もご覧ください。
Map(operands..., computation)
引数 | タイプ | セマンティクス |
---|---|---|
operands |
N 個の XlaOp のシーケンス |
N 個の型 T0..T{N-1} の配列 |
computation |
XlaComputation |
T 型の N 個のパラメータと任意の型の M を持つ T_0, T_1, .., T_{N + M -1} -> S 型の計算 |
dimensions |
int64 配列 |
地図のディメンションの配列 |
指定された operands
配列にスカラー関数を適用し、同じ次元の配列を生成します。各要素は、入力配列内の対応する要素に適用されたマッピングされた関数の結果です。
マッピングされた関数は任意の計算であり、スカラー型 T
の入力が N 個、S
型の出力が 1 個であるという制限があります。出力の次元は、要素の型 T が S で置き換えられている点を除き、オペランドと同じ次元になります。
たとえば、Map(op1, op2, op3, computation, par1)
は入力配列の各(多次元)インデックスで elem_out <-
computation(elem1, elem2, elem3, par1)
をマッピングして、出力配列を生成します。
OptimizationBarrier
最適化パスをブロックして、計算がバリアを越えて移動されないようにします。
バリアの出力に依存する演算子の前にすべての入力が評価されるようにします。
パッド
XlaBuilder::Pad
もご覧ください。
Pad(operand, padding_value, padding_config)
引数 | タイプ | セマンティクス |
---|---|---|
operand |
XlaOp |
T 型の配列 |
padding_value |
XlaOp |
追加されたパディングを埋める T 型のスカラー |
padding_config |
PaddingConfig |
両端(低、高)と各寸法の要素間のパディング量 |
配列の周囲と、指定された padding_value
を持つ配列の要素の間にパディングを追加して、指定された operand
配列を展開します。padding_config
は、各ディメンションのエッジ パディングと内部パディングを指定します。
PaddingConfig
は PaddingConfigDimension
の繰り返しフィールドで、ディメンションごとに 3 つのフィールド(edge_padding_low
、edge_padding_high
、interior_padding
)が含まれています。
edge_padding_low
と edge_padding_high
は、各ディメンションのローエンド(インデックス 0 の次)とハイエンド(最も高いインデックスの次)に追加するパディングの量をそれぞれ指定します。端のパディングの量は負の値でもかまいません。負のパディングの絶対値は、指定した寸法から削除する要素の数を示します。
interior_padding
は、各寸法の 2 つの要素間に追加されるパディングの量を指定します。負の値にすることはできません。内部パディングはエッジ パディングの前に論理的に発生するため、負エッジ パディングの場合、要素は内部パディング オペランドから削除されます。
エッジのパディングのペアがすべて(0, 0)で、内側のパディングの値がすべて 0 の場合、この演算は何も行われません。下の図は、2 次元配列の edge_padding
と interior_padding
のさまざまな値の例を示しています。
受信
XlaBuilder::Recv
もご覧ください。
Recv(shape, channel_handle)
引数 | タイプ | セマンティクス |
---|---|---|
shape |
Shape |
受信するデータの形状と、 |
channel_handle |
ChannelHandle |
送信と受信のペアごとに一意の識別子 |
同じチャネル ハンドルを共有する別の計算の Send
命令から、指定されたシェイプのデータを受け取ります。受信したデータの XLAOp を返します。
Recv
オペレーションのクライアント API は同期通信を表します。ただし、この命令は、非同期データ転送を可能にするために、内部で 2 つの HLO 命令(Recv
と RecvDone
)に分解されています。HloInstruction::CreateRecv
と HloInstruction::CreateRecvDone
もご覧ください。
Recv(const Shape& shape, int64 channel_id)
同じ channel_id を持つ Send
命令からデータを受信するために必要なリソースを割り当てます。割り当てられたリソースのコンテキストを返します。コンテキストは、次の RecvDone
命令で使用され、データ転送の完了を待機します。コンテキストは {受信バッファ(シェイプ)、リクエスト識別子(U32)} のタプルであり、RecvDone
命令でのみ使用できます。
RecvDone(HloInstruction context)
Recv
命令によって作成されたコンテキストが与えられ、データ転送が完了するまで待機し、受信したデータを返します。
削減
XlaBuilder::Reduce
もご覧ください。
1 つ以上の配列にリダクション関数を並行して適用します。
Reduce(operands..., init_values..., computation, dimensions)
引数 | タイプ | セマンティクス |
---|---|---|
operands |
N XlaOp のシーケンス |
T_0, ..., T_{N-1} 型の配列の N 個。 |
init_values |
N XlaOp のシーケンス |
T_0, ..., T_{N-1} 型の N 個のスカラー。 |
computation |
XlaComputation |
T_0, ..., T_{N-1}, T_0, ..., T_{N-1} -> Collate(T_0, ..., T_{N-1}) 型の計算。 |
dimensions |
int64 配列 |
削減するディメンションの配列。 |
ここで
- N は 1 以上である必要があります。
- 計算は「おおまかに」結合的である必要があります(以下を参照)。
- すべての入力配列のディメンションは同じである必要があります。
- すべての初期値は、
computation
の下で ID を形成する必要があります。 N = 1
の場合、Collate(T)
はT
です。N > 1
の場合、Collate(T_0, ..., T_{N-1})
はT
型のN
要素のタプルです。
この演算は、各入力配列の 1 つ以上の次元をスカラーに減らします。返される各配列のランクは rank(operand) - len(dimensions)
です。演算の出力は Collate(Q_0, ..., Q_N)
です。ここで、Q_i
は T_i
型の配列です。その次元については後述します。
異なるバックエンドで削減計算を再度関連付けることができます。加算などの一部のリダクション関数は浮動小数点数に対して結合的ではないため、数値の相違が生じる可能性があります。ただし、データの範囲が限られている場合は、ほとんどの実用的な用途において、浮動小数点数の加算で十分に結合的になります。
例
値 [10, 11,
12, 13]
とリダクション関数 f
(これは computation
)を持つ単一の 1 次元配列の 1 次元でリダクションする場合、次のように計算できます。
f(10, f(11, f(12, f(init_value, 13)))
他にも多くの可能性があります。
f(init_value, f(f(10, f(init_value, 11)), f(f(init_value, 12), f(init_value, 13))))
以下に、リダクションの実装方法の大まかな疑似コード例を示します。初期値を 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]
2 次元配列(行列)を縮小する例を示します。このシェイプはランク 2、サイズ 2 のディメンション 0、サイズ 3 のディメンション 1 です。
「add」関数でディメンション 0 または 1 を縮小した結果:
なお、リダクション結果はどちらも 1 次元配列です。この図では、見やすくするために列と行を示しています。
より複雑な例として、3 次元配列を示します。そのランクは 3、サイズ 4 のディメンション 0、サイズ 2 のディメンション 1、サイズ 3 のディメンション 2 です。わかりやすくするため、値 1 ~ 6 がディメンション 0 に複製されます。
2D の例と同様に、1 つの次元だけを縮小できます。たとえば、次元 0 を縮小すると、次元 0 のすべての値がスカラーに折りたたまれたランク 2 配列になります。
| 4 8 12 |
| 16 20 24 |
次元 2 を減らすと、次元 2 のすべての値がスカラーに折りたたまれたランク 2 配列も得られます。
| 6 15 |
| 6 15 |
| 6 15 |
| 6 15 |
入力の残りのディメンション間の相対的な順序は出力で保持されますが、一部のディメンションには(ランクが変わるため)新しい数値が割り当てられることがあります。
複数の次元を削減することもできます。次元 0 と 1 を加算減算すると、1 次元配列 [20, 28, 36]
が生成されます。
3 次元配列をそのすべての次元にわたって縮小すると、スカラー 84
が生成されます。
Variadic Reduce
N > 1
の場合、関数の削減の適用はすべての入力に同時に適用されるため、やや複雑になります。オペランドは次の順序で計算に渡されます。
- 第 1 オペランドの値を減らして実行する
- ...
- N 番目のオペランドの値を減らして実行する
- 第 1 オペランドの入力値
- ...
- N 番目のオペランドの入力値
たとえば、次のリダクション関数を使用すると、1 次元配列の max と argmax を並行して計算できます。
f: (Float, Int, Float, Int) -> Float, Int
f(max, argmax, value, index):
if value >= max:
return (value, index)
else:
return (max, argmax)
1 次元入力配列 V = Float[N], K = Int[N]
と init 値 I_V = Float, I_K = Int
の場合、唯一の入力ディメンション全体で削減の結果 f_(N-1)
は、次の再帰適用と同等です。
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))
このリダクションを値の配列と連続インデックスの配列(iota)に適用すると、配列全体で反復処理が行われ、最大値と一致するインデックスを含むタプルが返されます。
ReducePrecision
XlaBuilder::ReducePrecision
もご覧ください。
浮動小数点値を低精度形式(IEEE-FP16 など)に変換して元の形式に戻す場合の影響をモデル化します。低精度形式の指数と仮数のビット数は任意に指定できますが、すべてのビットサイズがすべてのハードウェア実装でサポートされているとは限りません。
ReducePrecision(operand, mantissa_bits, exponent_bits)
引数 | タイプ | セマンティクス |
---|---|---|
operand |
XlaOp |
浮動小数点型 T の配列。 |
exponent_bits |
int32 |
低精度形式の指数ビット数 |
mantissa_bits |
int32 |
低精度形式の仮数のビット数 |
結果は T
型の配列になります。入力値は、指定された仮数部のビット数で表現できる最も近い値に丸められ(「偶数と等しい」セマンティクスを使用)、指数ビット数で指定された範囲を超える値は正または負の無限大にクランプされます。NaN
値は保持されますが、正規の NaN
値に変換される場合があります。
低精度形式には、指数ビットが少なくとも 1 つ(どちらもゼロの仮数を持つため、ゼロ値と無限大を区別するために)必要です。また、仮数部のビット数が負でない必要があります。指数ビットまたは仮数ビットの数は、T
型の対応する値を超えることがあります。この場合、変換の対応する部分は何も行いません。
ReduceScatter
XlaBuilder::ReduceScatter
もご覧ください。
ReduceScatter は、AllReduce を効果的に行い、結果を scatter_dimension
に沿って shard_count
ブロックに分割し、ith
シャードを受け取るレプリカ グループ内のレプリカ i
に結果を分散させる集合的オペレーションです。
ReduceScatter(operand, computation, scatter_dim, shard_count,
replica_group_ids, channel_id)
引数 | タイプ | セマンティクス |
---|---|---|
operand |
XlaOp |
レプリカ全体で削減する配列または空でないタプル。 |
computation |
XlaComputation |
削減の計算 |
scatter_dimension |
int64 |
散布するディメンション。 |
shard_count |
int64 |
分割するブロック数: scatter_dimension |
replica_groups |
int64 のベクトルのベクトル |
リダクションが行われるグループ |
channel_id |
省略可 int64 |
モジュール間通信用のオプションのチャンネル ID |
operand
が配列のタプルの場合、タプルの各要素に対して Reduce-scatter が実行されます。replica_groups
は、リダクションが行われるレプリカ グループのリストです(現在のレプリカのレプリカ ID は、ReplicaId
を使用して取得できます)。各グループのレプリカの順序によって、all-reduce の結果が分散される順序が決まります。replica_groups
は空(この場合、すべてのレプリカが 1 つのグループに属している)か、レプリカの数と同じ数の要素が含まれている必要があります。レプリカ グループが複数ある場合は、すべて同じサイズにする必要があります。たとえば、replica_groups = {0, 2}, {1, 3}
はレプリカ0
と2
、および1
と3
の間でリダクションを実行し、結果を分散します。shard_count
は、各レプリカ グループのサイズです。これは、replica_groups
が空の場合に必要になります。replica_groups
が空でない場合、shard_count
は各レプリカ グループのサイズと同じにする必要があります。channel_id
はモジュール間通信に使用されます。互いに通信できるのは、同じchannel_id
を持つreduce-scatter
オペレーションのみです。
出力シェイプは、scatter_dimension
で shard_count
サイズに縮小された入力シェイプです。たとえば、2 つのレプリカがあり、2 つのレプリカでオペランドの値がそれぞれ [1.0, 2.25]
と [3.0, 5.25]
である場合、scatter_dim
が 0
であるこの演算の出力値は、最初のレプリカでは [4.0]
になり、2 番目のレプリカでは [7.5]
になります。
ReduceWindow
XlaBuilder::ReduceWindow
もご覧ください。
N 個の多次元配列のシーケンスの各ウィンドウ内のすべての要素にリダクション関数を適用し、N 個の多次元配列の単一またはタプルを出力として生成します。各出力配列には、ウィンドウの有効な位置の数と同じ数の要素が含まれます。プーリング レイヤは ReduceWindow
として表現できます。Reduce
と同様に、適用される computation
は常に左側の init_values
に渡されます。
ReduceWindow(operands..., init_values..., computation, window_dimensions,
window_strides, padding)
引数 | タイプ | セマンティクス |
---|---|---|
operands |
N XlaOps |
T_0,..., T_{N-1} 型の N 個の多次元配列のシーケンス。それぞれがウィンドウが配置される基本領域を表します。 |
init_values |
N XlaOps |
リダクションの開始値 N 個(N 個のオペランドごとに 1 個)。詳しくは、Reduce をご覧ください。 |
computation |
XlaComputation |
すべての入力オペランドの各ウィンドウ内の要素に適用する T_0, ..., T_{N-1}, T_0, ..., T_{N-1} -> Collate(T_0, ..., T_{N-1}) 型のリダクション関数。 |
window_dimensions |
ArraySlice<int64> |
ウィンドウ ディメンション値の整数の配列 |
window_strides |
ArraySlice<int64> |
ウィンドウ ストライド値の整数の配列 |
base_dilations |
ArraySlice<int64> |
基本拡張値の整数の配列 |
window_dilations |
ArraySlice<int64> |
ウィンドウ拡張値の整数の配列 |
padding |
Padding |
ウィンドウのパディング タイプ(Padding::kSame は、ストライドが 1 の場合に入力と同じ出力形状になるようにパディングします。または Padding::kValid はパディングを使用せず、収まらないとウィンドウを「停止」します) |
ここで
- N は 1 以上である必要があります。
- すべての入力配列のディメンションは同じである必要があります。
N = 1
の場合、Collate(T)
はT
です。N > 1
の場合、Collate(T_0, ..., T_{N-1})
は(T0,...T{N-1})
型のN
要素のタプルです。
次のコードと図に、ReduceWindow
の使用例を示します。入力はサイズ [4x6] の行列で、window_dimensions と window_stride_dimensions はどちらも [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);
ディメンションのストライドが 1 の場合、そのディメンション内のウィンドウの位置が、隣接するウィンドウから 1 要素離れていることを指定します。ウィンドウが互いに重複しないようにするには、window_stride_dimensions を window_dimensions と等しくする必要があります。以下の図は、2 つの異なるストライド値の使用を示しています。パディングは入力の各次元に適用され、計算は入力のパディングが加えられた場合と同じになります。
重要なパディングの例では、ディメンション 3
と入力配列 [10000, 1000, 100, 10, 1]
に対するストライド 2
で、Reduce-window の最小値(初期値は MAX_FLOAT
)を計算することを検討してください。パディング kValid
は、2 つの有効なウィンドウ([10000, 1000, 100]
と [100, 10, 1]
)で最小値を計算し、出力 [100, 1]
になります。まず kSame
をパディングすることで、配列の両側に初期要素を追加することで、Reduce-window の後の形状がストライド 1 の入力と同じになるように配列がパディングされます。これにより、[MAX_VALUE, 10000, 1000, 100, 10, 1,
MAX_VALUE]
になります。パディングされた配列に対して reduce-window を実行すると、[MAX_VALUE, 10000, 1000]
、[1000, 100, 10]
、[10, 1, MAX_VALUE]
の 3 つのウィンドウで動作し、[1000, 10, 1]
が生成されます。
リダクション関数の評価順序は任意であり、非決定的である場合もあります。したがって、リダクション関数が再関連付けの影響を受けすぎないようにする必要があります。詳細については、Reduce
のコンテキストでの関連付けに関する説明をご覧ください。
ReplicaId
XlaBuilder::ReplicaId
もご覧ください。
レプリカの一意の ID(U32 スカラー)を返します。
ReplicaId()
各レプリカの一意の ID は、[0, N)
区間の符号なし整数です。ここで、N
はレプリカの数です。すべてのレプリカが同じプログラムを実行しているため、プログラムの ReplicaId()
を呼び出すと、レプリカごとに異なる値が返されます。
Reshape
XlaBuilder::Reshape
オペレーションと Collapse
オペレーションもご覧ください。
配列の次元を新しい構成に再形成します。
Reshape(operand, new_sizes)
Reshape(operand, dimensions, new_sizes)
引数 | タイプ | セマンティクス |
---|---|---|
operand |
XlaOp |
T 型の配列 |
dimensions |
int64 個のベクトル |
ディメンションが折りたたまれた順序 |
new_sizes |
int64 個のベクトル |
新しい次元のサイズのベクトル |
概念的には、形状変更とは、まず配列をデータ値の 1 次元ベクトルにフラット化してから、このベクトルを新しい形状に改良することです。入力引数は、T 型の任意の配列、ディメンション インデックスのコンパイル時定数ベクトル、結果のディメンション サイズのコンパイル時定数ベクトルです。dimension
ベクトルの値を指定する場合は、すべての T 次元の順列にする必要があります。指定しない場合のデフォルトは {0, ..., rank - 1}
です。dimensions
のディメンションの順序は、ループネスト内で最も変化の遅いディメンション(最もメジャー)から最も変化の大きいディメンション(最もマイナー)で、入力配列が 1 つのディメンションに折りたたまれています。new_sizes
ベクトルによって、出力配列のサイズが決まります。new_sizes
のインデックス 0 の値はディメンション 0 のサイズ、インデックス 1 の値はディメンション 1 のサイズというようになります。new_size
ディメンションの積は、オペランドのディメンションサイズの積と等しくなければなりません。折りたたまれた配列を new_sizes
で定義された多次元配列に絞り込むと、new_sizes
のディメンションは変化の大きい順(最もメジャー)から最も変化が早い(最もマイナー)の順となります。
たとえば、v が 24 の要素からなる配列であるとします。
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} } };
特殊なケースとして、reshape で単一要素の配列をスカラーに変換できます。また、その逆も同様です。たとえば、
Reshape(f32[1x1] { {5} }, {0,1}, {}) == 5;
Reshape(5, {}, {1,1}) == f32[1x1] { {5} };
巻き戻し(反転)
XlaBuilder::Rev
もご覧ください。
Rev(operand, dimensions)
引数 | タイプ | セマンティクス |
---|---|---|
operand |
XlaOp |
T 型の配列 |
dimensions |
ArraySlice<int64> |
元に戻すディメンション |
指定された dimensions
に沿って operand
配列内の要素の順序を逆にして、同じ形状の出力配列を生成します。多次元インデックスにあるオペランド配列の各要素は、変換後のインデックスにある出力配列に格納されます。多次元インデックスは、各次元のインデックスを反転して変換されます。つまり、サイズ N の次元が逆次元の 1 つである場合、インデックス i は N - 1 - i に変換されます。
Rev
演算の用途の 1 つは、ニューラル ネットワークの勾配計算中に、2 つのウィンドウ次元に沿って畳み込み重み配列を反転することです。
RngNormal
XlaBuilder::RngNormal
もご覧ください。
正規分布に従って生成された乱数を使用して、特定の形状の出力を作成します。 \(N(\mu, \sigma)\) パラメータ \(\mu\) と \(\sigma\)、出力シェイプは浮動小数点の要素型である必要があります。さらに、パラメータはスカラー値である必要があります。
RngNormal(mu, sigma, shape)
引数 | タイプ | セマンティクス |
---|---|---|
mu |
XlaOp |
生成された数値の平均を指定する T 型のスカラー |
sigma |
XlaOp |
生成された値の標準偏差を指定する T 型のスカラー |
shape |
Shape |
T 型の出力シェイプ |
RngUniform
XlaBuilder::RngUniform
もご覧ください。
区間 \([a,b)\)の一様分布に従って生成された乱数を使用して、特定のシェイプの出力を作成します。パラメータと出力要素の型は、ブール型、整数型、または浮動小数点型でなければならず、型は一貫している必要があります。CPU バックエンドと GPU バックエンドは現在、F64、F32、F16、BF16、S64、U64、S32、U32 のみをサポートしています。また、パラメータはスカラー値である必要があります。 \(b <= a\) の場合、結果が実装で定義されます。
RngUniform(a, b, shape)
引数 | タイプ | セマンティクス |
---|---|---|
a |
XlaOp |
間隔の下限を指定する T 型のスカラー |
b |
XlaOp |
間隔の上限を指定する T 型のスカラー |
shape |
Shape |
T 型の出力シェイプ |
RngBitGenerator
指定されたアルゴリズム(またはバックエンドのデフォルト)を使用して、均一なランダムビットで満たされた特定の形状の出力を生成し、更新された状態(初期状態と同じ形状)と生成されたランダムデータを返します。
初期状態は、現在の乱数生成の初期状態です。必要な形状と有効な値は、使用するアルゴリズムによって異なります。
出力は、初期状態の決定論的関数であることが保証されますが、バックエンドと別のコンパイラ バージョンとの間で決定論的である保証はありません。
RngBitGenerator(algorithm, key, shape)
引数 | タイプ | セマンティクス |
---|---|---|
algorithm |
RandomAlgorithm |
使用する PRNG アルゴリズム。 |
initial_state |
XlaOp |
PRNG アルゴリズムの初期状態。 |
shape |
Shape |
生成されたデータの出力シェイプ。 |
algorithm
に使用可能な値:
rng_default
: バックエンド固有のシェイプ要件を持つバックエンド固有のアルゴリズム。rng_three_fry
: ThreeFry カウンタベースの PRNG アルゴリズム。initial_state
のシェイプはu64[2]
で、任意の値が設定されます。Salmon et al. SC 2011. 並行乱数: 1、2、3 のように簡単です。rng_philox
: 乱数を並列に生成する Philox アルゴリズム。initial_state
のシェイプは、任意の値のu64[3]
です。Salmon et al. SC 2011. 並行乱数: 1、2、3 のように簡単です。
散布
XLA 散布演算は、入力配列 operands
の値である一連の結果を生成します。複数のスライス(scatter_indices
で指定されたインデックスにある)は、update_computation
を使用して updates
の値のシーケンスで更新されます。
XlaBuilder::Scatter
もご覧ください。
scatter(operands..., scatter_indices, updates..., update_computation,
index_vector_dim, update_window_dims, inserted_window_dims,
scatter_dims_to_operand_dims)
引数 | タイプ | セマンティクス |
---|---|---|
operands |
N XlaOp のシーケンス |
分散する T_0, ..., T_N 型の N 個の配列。 |
scatter_indices |
XlaOp |
分散させる必要があるスライスの開始インデックスを含む配列。 |
updates |
N XlaOp のシーケンス |
T_0, ..., T_N 型の配列の N 個。updates[i] には、operands[i] の散布に必要な値が含まれています。 |
update_computation |
XlaComputation |
入力配列内の既存の値と分散中の更新を結合するために使用する計算。この計算は T_0, ..., T_N, T_0, ..., T_N -> Collate(T_0, ..., T_N) 型である必要があります。 |
index_vector_dim |
int64 |
開始インデックスを含む scatter_indices のディメンション。 |
update_window_dims |
ArraySlice<int64> |
ウィンドウの寸法である updates シェイプの寸法のセット。 |
inserted_window_dims |
ArraySlice<int64> |
updates シェイプに挿入する必要があるウィンドウのサイズのセット。 |
scatter_dims_to_operand_dims |
ArraySlice<int64> |
ディメンションは、散布インデックスからオペランド インデックス空間にマッピングされます。この配列は、i から scatter_dims_to_operand_dims[i] へのマッピングとして解釈されます。1 対 1 で合計にする必要があります。 |
indices_are_sorted |
bool |
インデックスが呼び出し元によって並べ替えられることが保証されているかどうか。 |
unique_indices |
bool |
インデックスが呼び出し元によって一意であることが保証されているかどうか。 |
ここで
- N は 1 以上である必要があります。
operands
[0
]、...、operands
[N-1
] はすべて同じディメンションである必要があります。updates
[0
]、...、updates
[N-1
] はすべて同じディメンションである必要があります。N = 1
の場合、Collate(T)
はT
です。N > 1
の場合、Collate(T_0, ..., T_N)
はT
型のN
要素のタプルです。
index_vector_dim
が scatter_indices.rank
と等しい場合、scatter_indices
の末尾に 1
ディメンションがあるものとして暗黙的に考慮します。
ArraySlice<int64>
型の update_scatter_dims
は、update_window_dims
にない updates
シェイプのディメンションのセットとして、昇順で定義されます。
scatter の引数は次の制約に従う必要があります。
各
updates
配列はupdate_window_dims.size + scatter_indices.rank - 1
ランクである必要があります。各
updates
配列のディメンションi
の境界は、次の条件を満たしている必要があります。i
がupdate_window_dims
に存在する(つまり、一部のk
についてupdate_window_dims
[k
] に等しい)場合、updates
のディメンションi
の境界は、inserted_window_dims
を考慮した後に、対応するoperand
の境界を超えてはなりません(adjusted_window_bounds
[k
]。adjusted_window_bounds
にはoperand
の境界が含まれており、インデックスinserted_window_dims
の境界は削除されています)。i
がupdate_scatter_dims
に存在する場合(つまり、一部のk
でupdate_scatter_dims
[k
] に等しい場合)、updates
のディメンションi
の境界は、対応するscatter_indices
の境界と等しくなければなりません。index_vector_dim
はスキップされます(k
<index_vector_dim
の場合はscatter_indices.shape.dims
[k
]、それ以外の場合はscatter_indices.shape.dims
[k+1
])。
update_window_dims
は昇順で、ディメンション番号を重複させないでください。また、[0, updates.rank)
の範囲内になければなりません。inserted_window_dims
は昇順で、ディメンション番号を重複させないでください。また、[0, operand.rank)
の範囲内になければなりません。operand.rank
は、update_window_dims.size
とinserted_window_dims.size
の合計と等しくなければなりません。scatter_dims_to_operand_dims.size
はscatter_indices.shape.dims
[index_vector_dim
] と等しくなければならず、その値は[0, operand.rank)
の範囲内になければなりません。
各 updates
配列内の特定のインデックス U
について、このアップデートを適用する必要がある、対応する operands
配列内の対応するインデックス I
は次のように計算されます。
G
= {U
[k
] forupdate_scatter_dims
inupdate_scatter_dims
} とします。G
を使用して、S
[i
] =scatter_indices
[Combine(G
,i
)] のように、scatter_indices
配列のインデックス ベクトルS
を検索します。ここで、Combine(A, b) は、b をindex_vector_dim
の位置で A に挿入します。k
scatter_dims_to_operand_dims
マップを使用してS
を分散し、S
を使用してoperand
にインデックスS
in
を作成します。正式名称は次のとおりです。S
k
<scatter_dims_to_operand_dims.size
の場合、in
[scatter_dims_to_operand_dims
[k
]] =S
[k
]。S
それ以外はin
[_
] =0
。
inserted_window_dims
に従ってU
のupdate_window_dims
にインデックスを分散し、各operands
配列にインデックスW
in
を作成します。正式名称は次のとおりです。W
in
[window_dims_to_operand_dims
(k
)] =U
[k
]k
がupdate_window_dims
の場合、window_dims_to_operand_dims
はドメイン [0
,update_window_dims.size
) と範囲 [0
,operand.rank
) \inserted_window_dims
を持つ単調関数です。(たとえば、update_window_dims.size
が4
、operand.rank
が6
、inserted_window_dims
が {0
、2
} の場合、window_dims_to_operand_dims
は {0
→1
、1
→3
、2
→4
、3
→5
} です)。W
それ以外はin
[_
] =0
。
I
はW
in
+S
in
です。「+」は要素ごとの加算です。
要約すると、分散オペレーションは次のように定義できます。
output
をoperands
で初期化します。つまり、operands
[J
] 配列内のすべてのインデックスO
について、すべてのインデックスJ
について初期化します。
output
[J
][O
] =operands
[J
][O
]updates
[J
] 配列のすべてのインデックスU
とoperand
[J
] 配列の対応するインデックスO
について、O
がoutput
の有効なインデックスである場合:
(output
[0
][O
], ...,output
[N-1
][O
]) =update_computation
(output
[0
][O
], ..., ,output
[N-1
][O
],updates
[0
][U
], ...,updates
[N-1
][U
])
更新が適用される順序は非決定的です。そのため、updates
内の複数のインデックスが operands
内の同じインデックスを参照している場合、output
内の対応する値は非決定的になります。
update_computation
に渡される最初のパラメータは常に output
配列の現在の値であり、2 番目のパラメータは常に updates
配列の値になります。これは、特に update_computation
が可換でない場合に重要です。
indices_are_sorted
が true に設定されている場合、XLA は start_indices
がユーザーによって(start_index_map
の昇順で)並べ替えられていると想定できます。そうでない場合、セマンティクスは実装が定義されています。
unique_indices
が true に設定されている場合、XLA は分散しているすべての要素が一意であると想定できます。そのため、XLA では非アトミック操作を使用できます。unique_indices
が true に設定され、分散しているインデックスが一意でない場合、セマンティクスは実装で定義されます。
非公式に、スキャッター オペレーションは収集オペレーションの逆と考えることができます。つまり、スキャッター オペレーションは、対応する収集オペレーションによって抽出された入力内の要素を更新します。
詳細な説明と例については、Gather
の「非公式の説明」セクションをご覧ください。
選択
XlaBuilder::Select
もご覧ください。
述語配列の値に基づいて、2 つの入力配列の要素から出力配列を作成します。
Select(pred, on_true, on_false)
引数 | タイプ | セマンティクス |
---|---|---|
pred |
XlaOp |
PRED 型の配列 |
on_true |
XlaOp |
T 型の配列 |
on_false |
XlaOp |
T 型の配列 |
配列 on_true
と on_false
の形状は同じである必要があります。これは出力配列の形状でもあります。配列 pred
は on_true
および on_false
と同じ次元で、要素型が PRED
である必要があります。
pred
の各要素 P
について、P
の値が true
の場合は on_true
から、P
の値が false
の場合は on_false
から、出力配列の対応する要素が取得されます。ブロードキャストの制限付き形式として、pred
は PRED
型のスカラーにできます。この場合、出力配列は、pred
が true
の場合は on_true
から、pred
が false
の場合は on_false
からすべて取得されます。
非スカラー pred
の例:
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};
スカラー pred
の例:
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};
タプル間の選択はサポートされています。この目的のため、タプルはスカラー型とみなされます。on_true
と on_false
がタプル(同じ形状である必要がある)の場合、pred
は PRED
型のスカラーである必要があります。
SelectAndScatter
XlaBuilder::SelectAndScatter
もご覧ください。
この演算は複合演算と考えることができます。まず operand
配列に対して ReduceWindow
を計算して各ウィンドウから要素を選択し、次に source
配列を選択された要素のインデックスに分散して、オペランド配列と同じ形状の出力配列を作成します。バイナリ select
関数は、各ウィンドウ全体に適用して各ウィンドウから要素を選択するために使用されます。この関数は、最初のパラメータのインデックス ベクトルが 2 番目のパラメータのインデックス ベクトルよりも辞書順で小さいというプロパティを使用して呼び出されます。select
関数は、最初のパラメータが選択された場合は true
を返し、2 番目のパラメータが選択された場合は false
を返します。また、関数は推移性を保持する必要があります(つまり、select(a, b)
と select(b, c)
が true
の場合、select(a, c)
も true
です)。これにより、選択された要素が特定のウィンドウで走査される要素の順序に依存しません。
関数 scatter
は、出力配列で選択された各インデックスに適用されます。次の 2 つのスカラー パラメータをとります。
- 出力配列で選択されたインデックスの現在の値
- 選択したインデックスに適用される
source
からの散布値。
2 つのパラメータを結合して、出力配列内の選択されたインデックスの値を更新するために使用されるスカラー値を返します。最初は、出力配列のすべてのインデックスが init_value
に設定されます。
出力配列は operand
配列と同じ形状になります。source
配列は、operand
配列に ReduceWindow
演算を適用した結果と同じ形状である必要があります。SelectAndScatter
を使用すると、ニューラル ネットワークのプーリングレイヤの勾配値を逆伝播できます。
SelectAndScatter(operand, select, window_dimensions, window_strides,
padding, source, init_value, scatter)
引数 | タイプ | セマンティクス |
---|---|---|
operand |
XlaOp |
ウィンドウがスライドする T 型の配列 |
select |
XlaComputation |
各ウィンドウ内のすべての要素に適用する T, T -> PRED 型のバイナリ計算。最初のパラメータが選択された場合は true を返し、2 番目のパラメータが選択された場合は false を返します。 |
window_dimensions |
ArraySlice<int64> |
ウィンドウ ディメンション値の整数の配列 |
window_strides |
ArraySlice<int64> |
ウィンドウ ストライド値の整数の配列 |
padding |
Padding |
ウィンドウのパディング タイプ(Padding::kSame または Padding::kValid) |
source |
XlaOp |
散布する値を含む T 型の配列 |
init_value |
XlaOp |
出力配列の初期値 T 型のスカラー値 |
scatter |
XlaComputation |
T, T -> T 型のバイナリ計算。散布図の各ソース要素をデスティネーション要素に適用します。 |
下の図は SelectAndScatter
の使用例を示しています。select
関数はパラメータの最大値を計算します。下の図(2)のようにウィンドウが重なっている場合、operand
配列のインデックスは異なるウィンドウで複数回選択される可能性があります。この図では、値 9 の要素がトップウィンドウ(青と赤)の両方によって選択され、バイナリ加算 scatter
関数によって値 8(2 + 6)の出力要素が生成されます。
scatter
関数の評価順序は任意であり、非決定的な場合もあります。そのため、scatter
関数が再関連付けの影響を受けすぎないようにする必要があります。詳細については、Reduce
のコンテキストでの関連付けに関する説明をご覧ください。
送信
XlaBuilder::Send
もご覧ください。
Send(operand, channel_handle)
引数 | タイプ | セマンティクス |
---|---|---|
operand |
XlaOp |
送信するデータ(T 型の配列) |
channel_handle |
ChannelHandle |
送信と受信のペアごとに一意の識別子 |
同じチャネル ハンドルを共有する別の計算の Recv
命令に、指定されたオペランド データを送信します。データは返されません。
Recv
オペレーションと同様に、Send
オペレーションのクライアント API は同期通信を表し、非同期データ転送を可能にするために内部で 2 つの HLO 命令(Send
と SendDone
)に分解されます。HloInstruction::CreateSend
と HloInstruction::CreateSendDone
もご覧ください。
Send(HloInstruction operand, int64 channel_id)
同じチャネル ID の Recv
命令によって割り当てられたリソースへのオペランドの非同期転送を開始します。コンテキストを返します。このコンテキストは、次の SendDone
命令によって使用され、データ転送の完了を待機します。コンテキストは {オペランド(シェイプ)、リクエスト識別子(U32)} のタプルであり、SendDone
命令でのみ使用できます。
SendDone(HloInstruction context)
Send
命令によって作成されたコンテキストが与えられ、データ転送が完了するまで待ちます。この命令はデータを返しません。
チャンネルに関する指示のスケジュール設定
各チャネル(Recv
、RecvDone
、Send
、SendDone
)に対する 4 つの命令の実行順序は次のとおりです。
Send
より前にRecv
が発生RecvDone
より前にSend
が発生RecvDone
より前にRecv
が発生SendDone
より前にSend
が発生
バックエンド コンパイラがチャネル命令を介して通信する計算ごとに線形スケジュールを生成する場合、計算間でサイクルがあってはなりません。たとえば、以下のスケジュールはデッドロックにつながります。
スライス
XlaBuilder::Slice
もご覧ください。
スライスでは、入力配列からサブ配列が抽出されます。サブ配列は入力と同じランクであり、入力配列内の境界ボックス内の値を含みます。ここで、境界ボックスのディメンションとインデックスは、スライス演算への引数として指定されます。
Slice(operand, start_indices, limit_indices, strides)
引数 | タイプ | セマンティクス |
---|---|---|
operand |
XlaOp |
T 型の N 次元配列 |
start_indices |
ArraySlice<int64> |
各ディメンションのスライスの開始インデックスを含む N 個の整数のリスト。0 以上の値を指定してください。 |
limit_indices |
ArraySlice<int64> |
各ディメンションのスライスの終了インデックス(含まれない)を含む N 個の整数のリスト。各値は、ディメンションのそれぞれの start_indices 値以上、ディメンションのサイズ以下にする必要があります。 |
strides |
ArraySlice<int64> |
スライスの入力ストライドを決定する N 個の整数のリスト。スライスは、ディメンション d のすべての strides[d] 要素を選択します。 |
1 次元の例:
let a = {0.0, 1.0, 2.0, 3.0, 4.0}
Slice(a, {2}, {4}) produces:
{2.0, 3.0}
2 次元の例:
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} }
並べ替え
XlaBuilder::Sort
もご覧ください。
Sort(operands, comparator, dimension, is_stable)
引数 | タイプ | セマンティクス |
---|---|---|
operands |
ArraySlice<XlaOp> |
並べ替えるオペランド。 |
comparator |
XlaComputation |
使用するコンパレータ計算。 |
dimension |
int64 |
並べ替えるディメンション。 |
is_stable |
bool |
安定した並べ替えを使用するかどうか。 |
オペランドを 1 つだけ指定した場合:
オペランドがランク 1 テンソル(配列)の場合、結果は並べ替えられた配列になります。配列を昇順で並べ替えるには、コンパレータで小なりの比較を行う必要があります。形式的には、配列の並べ替え後に、
i < j
を持つすべてのインデックス位置i, j
(comparator(value[i], value[j]) = comparator(value[j], value[i]) = false
またはcomparator(value[i], value[j]) = true
)が保持されます。オペランドの方がランクが高い場合、オペランドは指定されたディメンションに沿って並べ替えられます。たとえば、階数 2 テンソル(行列)の場合、ディメンション値
0
はすべての列を独立して並べ替え、ディメンション値1
は各行を個別に並べ替えます。ディメンション番号を指定しない場合、デフォルトで最後のディメンションが選択されます。並べ替えられるディメンションには、ランク 1 の場合と同じ並べ替え順序が適用されます。
n > 1
オペランドを指定した場合:
すべての
n
オペランドは同じ次元のテンソルでなければなりません。テンソルの要素の型は異なる場合があります。すべてのオペランドは個別にではなく、一緒に並べ替えられます。概念的には、オペランドはタプルとして扱われます。インデックス位置
i
とj
にある各オペランドの要素を入れ替える必要があるかどうかをチェックする場合は、2 * n
スカラー パラメータを使用してコンパレータが呼び出されます。ここで、パラメータ2 * k
はk-th
オペランドの位置i
の値に対応し、パラメータ2 * k + 1
はk-th
オペランドの位置j
の値に対応します。したがって、通常、コンパレータはパラメータ2 * k
と2 * k + 1
を相互に比較し、場合によっては他のパラメータペアをタイブレーカーとして使用します。結果は、並べ替えられた順序で(上記のように指定されたディメンションに沿って)オペランドで構成されるタプルになります。タプルの
i-th
オペランドは、Sort のi-th
オペランドに対応します。
たとえば、3 つのオペランド operand0 = [3, 1]
、operand1 = [42, 50]
、operand2 = [-3.0, 1.1]
があり、コンパレータが operand0
の値のみと「より小さい」を比較した場合、並べ替えの出力はタプル ([1, 3], [50, 42], [1.1, -3.0])
になります。
is_stable
が true に設定されている場合、並べ替えが安定していることが保証されます。つまり、コンパレータによって等しいとみなされる要素がある場合、等しい値の相対的な順序が保持されます。e1
と e2
の 2 つの要素は、comparator(e1, e2) = comparator(e2, e1) = false
の場合のみ等しくなります。デフォルトでは、is_stable
は false に設定されています。
行 / 列の入れ替え
tf.reshape
オペレーションもご覧ください。
Transpose(operand)
引数 | タイプ | セマンティクス |
---|---|---|
operand |
XlaOp |
転置するオペランド。 |
permutation |
ArraySlice<int64> |
ディメンションを並べ替える方法。 |
指定された順列でオペランド次元を並べ替えます。したがって、∀ i . 0 ≤ i < rank ⇒ input_dimensions[permutation[i]] = output_dimensions[i]
。
これは Reshape(operand, permutation, Permute(permutation, operand.shape.dimensions)) と同じです。
TriangularSolve
XlaBuilder::TriangularSolve
もご覧ください。
前方置換または逆置換によって、下三角係数行列または上三角係数行列を持つ連立一次方程式を解きます。先行次元に沿ってブロードキャストし、このルーティンは変数 x
について、行列系 op(a) * x =
b
または x * op(a) = b
のいずれかを解きます。指定された a
と b
(op(a)
は op(a) = a
、op(a) = Transpose(a)
、op(a) = Conj(Transpose(a))
のいずれか)です。
TriangularSolve(a, b, left_side, lower, unit_diagonal, transpose_a)
引数 | タイプ | セマンティクス |
---|---|---|
a |
XlaOp |
形状が [..., M, M] の複素数または浮動小数点型のランク > 2 配列。 |
b |
XlaOp |
left_side が true の場合は [..., M, K] の形状で、同じ型のランク > 2 の配列。それ以外の場合は [..., K, M] 。 |
left_side |
bool |
op(a) * x = b (true )または x * op(a) = b (false )の形式の系を解くかどうかを示します。 |
lower |
bool |
a の上または下の三角形を使用するかどうか。 |
unit_diagonal |
bool |
true の場合、a の対角要素は 1 とみなされ、アクセスされません。 |
transpose_a |
Transpose |
a をそのまま使用するか、転置するか、共役の転置を取るかを指定します。 |
入力データは、lower
の値に応じて、a
の下/上三角形からのみ読み取られます。もう一方の三角形の値は無視されます。出力データは同じ三角形内に返されます。もう一方の三角形の値は実装で定義されるもので、任意の値を指定できます。
a
と b
のランクが 2 より大きい場合、これらは行列のバッチとして扱われ、マイナー 2 次元を除くすべての次元がバッチ ディメンションです。a
と b
のバッチ ディメンションは等しい必要があります。
タプル
XlaBuilder::Tuple
もご覧ください。
可変数のデータハンドルを含むタプル。各ハンドルは独自の形状を持ちます。
これは C++ の std::tuple
に似ています。概念的には次のようになります。
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);
タプルは、GetTupleElement
オペレーションで分解(アクセス)できます。
一方
XlaBuilder::While
もご覧ください。
While(condition, body, init)
引数 | タイプ | セマンティクス |
---|---|---|
condition |
XlaComputation |
ループの終了条件を定義する T -> PRED 型の XlaComputation。 |
body |
XlaComputation |
ループの本文を定義する T -> T 型の XlaComputation。 |
init |
T |
condition と body のパラメータの初期値。 |
condition
が失敗するまで、body
を順次実行します。これは、以下の相違点と制限事項を除き、他の多くの言語の通常の while ループと似ています。
While
ノードはT
型の値を返します。これは、body
を最後に実行した結果です。T
型の形状は静的に決定され、すべての反復処理で同一である必要があります。
計算の T パラメータは、最初の反復処理で init
値で初期化され、後続の反復処理ごとに body
の新しい結果に自動的に更新されます。
While
ノードの主なユースケースの 1 つは、ニューラル ネットワークでのトレーニングの繰り返し実行を実装することです。簡略化した疑似コードと、計算を表すグラフを以下に示します。コードは while_test.cc
にあります。この例の T
型は、反復回数の int32
とアキュムレータの vector[10]
で構成される Tuple
です。1, 000 回のイテレーションで、ループは定数ベクトルをアキュムレータに追加し続けます。
// 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};
}