オペレーション セマンティクス

以下では、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} はレプリカ 02 と、13 を連結します。
  • shard_count は、各レプリカ グループのサイズです。これは、replica_groups が空の場合に必要になります。
  • channel_id はモジュール間通信に使用されます。互いに通信できるのは、同じ channel_id を持つ all-gather オペレーションのみです。

出力シェイプは、all_gather_dimshard_count 倍になった入力シェイプです。たとえば、2 つのレプリカがあり、2 つのレプリカでオペランドの値が [1.0, 2.5][3.0, 5.25] の場合、all_gather_dim0 であるこの演算の出力値は両方のレプリカで [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} はレプリカ 02 と、13 の間でリダクションを行います。
  • 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 つのフェーズがあります。

  1. 分散フェーズ。各コアで、オペランドは split_dimensions に沿って split_count 個のブロックに分割され、ブロックはすべてのコアに分散されます。たとえば、i 番目のブロックは i 番目のコアに送信されます。
  2. 収集フェーズ。各コアは、受信したブロックを 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_indexoperand の特徴ディメンションのインデックス)ごとに、他のすべてのディメンションの operandoffsetscale に関する勾配を計算します。feature_index は、operand 内の特徴ディメンションの有効なインデックスである必要があります。

3 つの勾配は、次の式で定義されます(4 次元配列を operand、特徴ディメンションのインデックス l、バッチサイズ m、空間サイズ wh を持つと仮定)。

\[ \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} \]

入力 meanvariance は、バッチ ディメンションと空間ディメンション全体でのモーメント値を表します。

出力タイプは、次の 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_indexoperand の特徴ディメンションのインデックス)に対して、他のすべてのディメンションの平均と分散を計算し、平均と分散を使用して operand の各要素を正規化します。feature_index には、operand 内の特徴ディメンションの有効なインデックスを指定する必要があります。

BatchNormInference は、バッチごとに meanvariance を計算せずに BatchNormTraining を呼び出す場合と同じです。代わりに、入力 meanvariance を推定値として使用します。この op の目的は推論のレイテンシを短縮することであるため、BatchNormInference という名前が付けられています。

出力は、入力 operand と同じ形状の N 次元の正規化配列です。

BatchNormTraining

アルゴリズムの詳細な説明については、XlaBuilder::BatchNormTrainingthe 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_indexoperand の特徴ディメンションのインデックス)に対して、他のすべてのディメンションの平均と分散を計算し、平均と分散を使用して operand の各要素を正規化します。feature_index には、operand 内の特徴ディメンションの有効なインデックスを指定する必要があります。

空間次元のサイズとして wh を持つ 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_meanbatch_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 の上または下の三角形を使用するかどうか。

lowertrue の場合、$a = l となるように下三角行列 l を計算します。 l^T$。lowerfalse の場合、\(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 つの配列はすべて同じ形状である必要があります。また、ブロードキャストの制限付き形式として、minmaxT 型のスカラーにすることもできます。

スカラー minmax の例:

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

predtrue の場合は true_computationpredfalse の場合は false_computation を実行し、結果を返します。

true_computation は、 \(T_0\) 型の引数を 1 つ受け取る必要があり、同じ型である true_operand で呼び出されます。false_computation は、 \(T_1\) 型の引数を 1 つ受け取る必要があり、同じ型である false_operand で呼び出されます。true_computationfalse_computation の戻り値の型は同じである必要があります。

pred の値に応じて、true_computationfalse_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 次元畳み込みについて考え、出力に固定された batchzyx 座標を選択します。この場合、(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 つ目の次元の積の合計を実行します。これらは「契約」ディメンションです。lhsrhs の契約サイズは同じサイズにする必要があります。実際には、ベクトル間のドット積、ベクトル/行列の乗算、行列/行列の乗算に使用できます。

DotGeneral

XlaBuilder::DotGeneral もご覧ください。

DotGeneral(lhs, rhs, dimension_numbers)

引数 タイプ セマンティクス
lhs XlaOp T 型の配列
rhs XlaOp T 型の配列
dimension_numbers DotDimensionNumbers 分割とバッチのディメンション番号

Dot と似ていますが、lhsrhs の両方に縮小ディメンションとバッチ ディメンション番号を指定できます。

DotDimensionNumbers フィールド タイプ セマンティクス
lhs_contracting_dimensions 繰り返し int64 lhs 契約ディメンション番号
rhs_contracting_dimensions 繰り返し int64 rhs 契約ディメンション番号
lhs_batch_dimensions 繰り返し int64 lhs 個のバッチ ディメンション番号
rhs_batch_dimensions 繰り返し int64 rhs 個のバッチ ディメンション番号

DotGeneral は、dimension_numbers で指定された分割ディメンションに対する積の合計を実行します。

lhsrhs の関連付けられた縮小ディメンション番号は同じにする必要はありませんが、同じディメンション サイズにする必要があります。

ディメンション番号を縮小する例:

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} }

lhsrhs の関連するバッチ ディメンション番号は、同じディメンション サイズである必要があります。

バッチ ディメンション番号の例(バッチサイズ 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 の値である結果を生成します。スライス updatestart_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)

ここで、OpAdd(加算)、Sub(減算)、Mul(乗算)、Div(除算)、Rem(剰余)、Max(最大)、Min(最小)、LogicalAnd(論理 AND)、LogicalOr(論理 OR)のいずれかです。

引数 タイプ セマンティクス
lhs XlaOp 左側のオペランド: T 型の配列
rhs XlaOp 右側のオペランド: T 型の配列

引数の形状は類似しているか、互換性がある必要があります。シェイプの互換性の意味については、ブロードキャストのドキュメントをご覧ください。演算の結果の形状は、2 つの入力配列をブロードキャストした結果です。このバリアントでは、オペランドの 1 つがスカラーでない限り、異なるランクの配列間の演算はサポートされません。

OpRem の場合、結果の符号は被除数から取得され、結果の絶対値は常に除数の絶対値よりも小さくなります。

整数除算オーバーフロー(ゼロによる符号付き/符号なし除算/剰余、または -1 を使用した INT_SMIN の符号付き除算/剰余)は、実装で定義された値を生成します。

次のオペレーションには、異なるランクのブロードキャストをサポートする別のバリアントがあります。

Op(lhs, rhs, broadcast_dimensions)

ここで、Op は上記と同じです。この演算のバリアントは、異なるランクの配列間の算術演算(ベクトルへの行列の追加など)に使用します。

追加の broadcast_dimensions オペランドは、低ランクオペランドのランクを高ランクオペランドのランクに拡張するために使用される整数のスライスです。broadcast_dimensions は、下位シェイプのディメンションを上位シェイプのディメンションにマッピングします。展開されたシェイプのマッピングされていないサイズは、サイズ 1 のディメンションで埋められます。次に、縮退次元ブロードキャストは、これらの縮退次元に沿って形状をブロードキャストして、両方のオペランドの形状を均等にします。セマンティクスの詳細については、ブロードキャストのページをご覧ください。

要素単位の比較演算

XlaBuilder::Eq もご覧ください。

標準的な要素単位のバイナリ比較演算のセットがサポートされています。浮動小数点型を比較する場合は、標準の IEEE 754 浮動小数点比較セマンティクスが適用されます。

Op(lhs, rhs)

ここで、OpEq(等しい)、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.sizecollapsed_slice_dims.size の合計と等しくなければなりません。また、slice_sizes.sizeoperand.rank と等しくする必要があります。

index_vector_dimstart_indices.rank と等しい場合、暗黙的に start_indices に末尾の 1 ディメンションがあると見なします(つまり、start_indices の形状が [6,7] で、index_vector_dim2 の場合、start_indices の形状は暗黙的に [6,7,1] であると見なされます)。

ディメンション i に沿った出力配列の境界は次のように計算されます。

  1. ibatch_dims に存在する(一部の kbatch_dims[k] に等しい)場合、対応するディメンションの境界が start_indices.shape から選択され、index_vector_dim はスキップされます(k < index_vector_dim の場合は start_indices.shape.dims[k]、それ以外の場合は start_indices.shape.dims[k+1] が選択されます)。

  2. ioffset_dims に存在する(一部の koffset_dims[k] に等しい)場合、collapsed_slice_dims を考慮したうえで、slice_sizes から対応する境界を選択します(つまり、adjusted_slice_sizesslice_sizes で、インデックス collapsed_slice_dims の境界を削除した adjusted_slice_sizes[k] を選択します)。

特定の出力インデックス Out に対応するオペランド インデックス In は、形式的に次のように計算されます。

  1. 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 となります。

  2. start_index_map を使用して S を分散し、S を使用して operand に開始インデックス Sin を作成します。より正確に言うと、

    1. Sk < start_index_map.size の場合、in[start_index_map[k]] = S[k]。

    2. Sそれ以外は in[_] = 0

  3. collapsed_slice_dims のセットに従って Out のオフセット ディメンションでインデックスを分散させることで、operand へのインデックス Oin を作成します。より正確に言うと、

    1. Ok < offset_dims.size の場合、in[remapped_offset_dims(k)] = Out[offset_dims[k]](remapped_offset_dims は以下で定義)。

    2. Oそれ以外は in[_] = 0

  4. InOin + Sin です。「+」は要素ごとの加算です。

remapped_offset_dims は、ドメイン [0, offset_dims.size)、範囲 [0, operand.rank) \ collapsed_slice_dims を持つ単調関数です。たとえばoffset_dims.size4operand.rank6collapsed_slice_dims は {02}、remapped_offset_dims は {01132435} です。

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_dimstart_indices.rank - 1 に設定されています。index_vector_dim の値に興味深い値を指定すると、演算自体は変わりませんが、視覚的表現が複雑になります。

上記すべてがどのように組み合わされているかを理解するために、[16,11] 配列からシェイプ [8,6] の 5 つのスライスを収集する例を見てみましょう。[16,11] 配列内のスライスの位置は、形状 S64[2] のインデックス ベクトルとして表すことができるため、5 つの位置のセットは S64[5,2] 配列として表すことができます。

収集オペレーションの動作は、出力シェイプのインデックスである [GO0O1] を取り、次のように入力配列の要素にマッピングするインデックス変換として表すことができます。

まず、G を使用して collect インデックス配列から(X,Y)ベクトルを選択します。出力配列のインデックス [GO0O1] にある要素は、インデックス [X+O0,Y+O1] にある入力配列の要素です。

slice_sizes[8,6] です。これにより、O0 と O1 の範囲が決定され、スライスの境界が決定されます。

この収集オペレーションは、G をバッチ ディメンションとするバッチ動的スライスとして機能します。

ギャザー インデックスは多次元の場合があります。たとえば、上記の例のより一般的なバージョンでシェイプ [4,5,2] の「gather indexs」配列を使用すると、次のようなインデックスが変換されます。

ここでも、バッチ動的スライス G0G1 をバッチ ディメンションとして機能します。スライスのサイズは [8,6] のままです。

XLA での収集オペレーションでは、上記で説明した非公式のセマンティクスが次のように一般化されます。

  1. 出力シェイプのどのディメンションがオフセット ディメンションかを構成できます(最後の例では、O0O1 を含むディメンション)。出力バッチ ディメンション(最後の例の G0G1 を含むディメンション)は、オフセット ディメンションではない出力ディメンションとして定義されます。

  2. 出力シェイプに明示的に存在する出力オフセット次元の数は、入力ランクよりも小さくなる場合があります。このような「欠落」しているディメンションは collapsed_slice_dims として明示的にリストされ、スライスサイズを 1 にする必要があります。スライスサイズが 1 であるため、有効なインデックスは 0 のみであり、省略してもあいまいさは発生しません。

  3. 「Collect Indices」配列(最後の例では XY)から抽出されたスライスの要素が入力配列ランクよりも少ない場合があり、明示的なマッピングにより、入力と同じランクになるようにインデックスを拡張する方法が指示されます。

最後の例として、(2)と(3)を使用して tf.gather_nd を実装します。

G0G1 は、開始インデックスに X という 1 つの要素しかないことを除き、通常と同様に collect インデックス配列から開始インデックスをスライスするために使用されます。同様に、値が O0 の出力オフセット インデックスは 1 つだけです。ただし、入力配列のインデックスとして使用される前に、これらは「集約インデックス マッピング」(正式な説明の start_index_map)と「オフセット マッピング」(正式な記述の remapped_offset_dims)に従ってそれぞれ [X0] と [0O0] に展開され、[XO0{2] の [G01 1] に拡張されます。つまり、出力インデックス [0GO] は [01] に拡張されます。つまり、[0GO] はインデックス [01] にマップされます。G1GatherIndicestf.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 は、各ディメンションのエッジ パディングと内部パディングを指定します。

PaddingConfigPaddingConfigDimension の繰り返しフィールドで、ディメンションごとに 3 つのフィールド(edge_padding_lowedge_padding_highinterior_padding)が含まれています。

edge_padding_lowedge_padding_high は、各ディメンションのローエンド(インデックス 0 の次)とハイエンド(最も高いインデックスの次)に追加するパディングの量をそれぞれ指定します。端のパディングの量は負の値でもかまいません。負のパディングの絶対値は、指定した寸法から削除する要素の数を示します。

interior_padding は、各寸法の 2 つの要素間に追加されるパディングの量を指定します。負の値にすることはできません。内部パディングはエッジ パディングの前に論理的に発生するため、負エッジ パディングの場合、要素は内部パディング オペランドから削除されます。

エッジのパディングのペアがすべて(0, 0)で、内側のパディングの値がすべて 0 の場合、この演算は何も行われません。下の図は、2 次元配列の edge_paddinginterior_padding のさまざまな値の例を示しています。

受信

XlaBuilder::Recv もご覧ください。

Recv(shape, channel_handle)

引数 タイプ セマンティクス
shape Shape 受信するデータの形状と、
channel_handle ChannelHandle 送信と受信のペアごとに一意の識別子

同じチャネル ハンドルを共有する別の計算の Send 命令から、指定されたシェイプのデータを受け取ります。受信したデータの XLAOp を返します。

Recv オペレーションのクライアント API は同期通信を表します。ただし、この命令は、非同期データ転送を可能にするために、内部で 2 つの HLO 命令(RecvRecvDone)に分解されています。HloInstruction::CreateRecvHloInstruction::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_iT_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} はレプリカ 02、および 13 の間でリダクションを実行し、結果を分散します。
  • shard_count は、各レプリカ グループのサイズです。これは、replica_groups が空の場合に必要になります。replica_groups が空でない場合、shard_count は各レプリカ グループのサイズと同じにする必要があります。
  • channel_id はモジュール間通信に使用されます。互いに通信できるのは、同じ channel_id を持つ reduce-scatter オペレーションのみです。

出力シェイプは、scatter_dimensionshard_count サイズに縮小された入力シェイプです。たとえば、2 つのレプリカがあり、2 つのレプリカでオペランドの値がそれぞれ [1.0, 2.25][3.0, 5.25] である場合、scatter_dim0 であるこの演算の出力値は、最初のレプリカでは [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 に使用可能な値:

散布

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_dimscatter_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 の境界は、次の条件を満たしている必要があります。

    • iupdate_window_dims に存在する(つまり、一部の k について update_window_dims[k] に等しい)場合、updates のディメンション i の境界は、inserted_window_dims を考慮した後に、対応する operand の境界を超えてはなりません(adjusted_window_bounds[k]。adjusted_window_bounds には operand の境界が含まれており、インデックス inserted_window_dims の境界は削除されています)。
    • iupdate_scatter_dims に存在する場合(つまり、一部の kupdate_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.sizeinserted_window_dims.size の合計と等しくなければなりません。

  • scatter_dims_to_operand_dims.sizescatter_indices.shape.dims[index_vector_dim] と等しくなければならず、その値は [0, operand.rank) の範囲内になければなりません。

updates 配列内の特定のインデックス U について、このアップデートを適用する必要がある、対応する operands 配列内の対応するインデックス I は次のように計算されます。

  1. G = { U[k] for update_scatter_dims in update_scatter_dims } とします。G を使用して、S[i] = scatter_indices[Combine(G, i)] のように、scatter_indices 配列のインデックス ベクトル S を検索します。ここで、Combine(A, b) は、b を index_vector_dim の位置で A に挿入します。k
  2. scatter_dims_to_operand_dims マップを使用して S を分散し、S を使用して operand にインデックス Sin を作成します。正式名称は次のとおりです。
    1. Sk < scatter_dims_to_operand_dims.size の場合、in[scatter_dims_to_operand_dims[k]] = S[k]。
    2. Sそれ以外は in[_] = 0
  3. inserted_window_dims に従って Uupdate_window_dims にインデックスを分散し、各 operands 配列にインデックス Win を作成します。正式名称は次のとおりです。
    1. Win[window_dims_to_operand_dims(k)] = U[k] kupdate_window_dims の場合、window_dims_to_operand_dims はドメイン [0, update_window_dims.size) と範囲 [0, operand.rank) \ inserted_window_dims を持つ単調関数です。(たとえば、update_window_dims.size4operand.rank6inserted_window_dims が {02} の場合、window_dims_to_operand_dims は {01132435} です)。
    2. Wそれ以外は in[_] = 0
  4. IWin + Sin です。「+」は要素ごとの加算です。

要約すると、分散オペレーションは次のように定義できます。

  • outputoperands で初期化します。つまり、operands[J] 配列内のすべてのインデックス O について、すべてのインデックス J について初期化します。
    output[J][O] = operands[J][O]
  • updates[J] 配列のすべてのインデックス Uoperand[J] 配列の対応するインデックス O について、Ooutput の有効なインデックスである場合:
    (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_trueon_false の形状は同じである必要があります。これは出力配列の形状でもあります。配列 predon_true および on_false と同じ次元で、要素型が PRED である必要があります。

pred の各要素 P について、P の値が true の場合は on_true から、P の値が false の場合は on_false から、出力配列の対応する要素が取得されます。ブロードキャストの制限付き形式として、predPRED 型のスカラーにできます。この場合、出力配列は、predtrue の場合は on_true から、predfalse の場合は 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_trueon_false がタプル(同じ形状である必要がある)の場合、predPRED 型のスカラーである必要があります。

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 つのスカラー パラメータをとります。

  1. 出力配列で選択されたインデックスの現在の値
  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 命令(SendSendDone)に分解されます。HloInstruction::CreateSendHloInstruction::CreateSendDone もご覧ください。

Send(HloInstruction operand, int64 channel_id)

同じチャネル ID の Recv 命令によって割り当てられたリソースへのオペランドの非同期転送を開始します。コンテキストを返します。このコンテキストは、次の SendDone 命令によって使用され、データ転送の完了を待機します。コンテキストは {オペランド(シェイプ)、リクエスト識別子(U32)} のタプルであり、SendDone 命令でのみ使用できます。

SendDone(HloInstruction context)

Send 命令によって作成されたコンテキストが与えられ、データ転送が完了するまで待ちます。この命令はデータを返しません。

チャンネルに関する指示のスケジュール設定

各チャネル(RecvRecvDoneSendSendDone)に対する 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, jcomparator(value[i], value[j]) = comparator(value[j], value[i]) = false または comparator(value[i], value[j]) = true)が保持されます。

  • オペランドの方がランクが高い場合、オペランドは指定されたディメンションに沿って並べ替えられます。たとえば、階数 2 テンソル(行列)の場合、ディメンション値 0 はすべての列を独立して並べ替え、ディメンション値 1 は各行を個別に並べ替えます。ディメンション番号を指定しない場合、デフォルトで最後のディメンションが選択されます。並べ替えられるディメンションには、ランク 1 の場合と同じ並べ替え順序が適用されます。

n > 1 オペランドを指定した場合:

  • すべての n オペランドは同じ次元のテンソルでなければなりません。テンソルの要素の型は異なる場合があります。

  • すべてのオペランドは個別にではなく、一緒に並べ替えられます。概念的には、オペランドはタプルとして扱われます。インデックス位置 ij にある各オペランドの要素を入れ替える必要があるかどうかをチェックする場合は、2 * n スカラー パラメータを使用してコンパレータが呼び出されます。ここで、パラメータ 2 * kk-th オペランドの位置 i の値に対応し、パラメータ 2 * k + 1k-th オペランドの位置 j の値に対応します。したがって、通常、コンパレータはパラメータ 2 * k2 * 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 に設定されている場合、並べ替えが安定していることが保証されます。つまり、コンパレータによって等しいとみなされる要素がある場合、等しい値の相対的な順序が保持されます。e1e2 の 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 のいずれかを解きます。指定された abop(a)op(a) = aop(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 = btrue)または x * op(a) = bfalse)の形式の系を解くかどうかを示します。
lower bool a の上または下の三角形を使用するかどうか。
unit_diagonal bool true の場合、a の対角要素は 1 とみなされ、アクセスされません。
transpose_a Transpose a をそのまま使用するか、転置するか、共役の転置を取るかを指定します。

入力データは、lower の値に応じて、a の下/上三角形からのみ読み取られます。もう一方の三角形の値は無視されます。出力データは同じ三角形内に返されます。もう一方の三角形の値は実装で定義されるもので、任意の値を指定できます。

ab のランクが 2 より大きい場合、これらは行列のバッチとして扱われ、マイナー 2 次元を除くすべての次元がバッチ ディメンションです。ab のバッチ ディメンションは等しい必要があります。

タプル

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 conditionbody のパラメータの初期値。

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};
}