このページは Cloud Translation API によって翻訳されました。

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

以下では、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` の特徴ディメンションのインデックス

空間次元のサイズとして 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];
  // ...
}

ユーザー指定関数に副作用があってはなりません。また、その実行はべき等でなければなりません。

注: ユーザー提供の関数は不透明な性質を持つため、コンパイラの最適化の機会が制限されます。可能な限り、ネイティブ XLA 演算で計算を表現するようにしてください。CustomCall は最後の手段としてのみ使用してください。

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 を作成します。より正確に言うと、
1. Sk < start_index_map.size の場合、_in[start_index_map[k]] = S[k]。
2. Sそれ以外は _in[_] = 0。
collapsed_slice_dims のセットに従って Out のオフセットディメンションでインデックスを分散させることで、operand へのインデックス O_in を作成します。より正確に言うと、
1. Ok < offset_dims.size の場合、_in[remapped_offset_dims(k)] = Out[offset_dims[k]]（remapped_offset_dims は以下で定義）。
2. 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₀、O₁] を取り、次のように入力配列の要素にマッピングするインデックス変換として表すことができます。

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

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

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

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

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

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

出力シェイプのどのディメンションがオフセットディメンションかを構成できます（最後の例では、O₀、O₁ を含むディメンション）。出力バッチディメンション（最後の例の G₀、G₁ を含むディメンション）は、オフセットディメンションではない出力ディメンションとして定義されます。
出力シェイプに明示的に存在する出力オフセット次元の数は、入力ランクよりも小さくなる場合があります。このような「欠落」しているディメンションは collapsed_slice_dims として明示的にリストされ、スライスサイズを 1 にする必要があります。スライスサイズが 1 であるため、有効なインデックスは 0 のみであり、省略してもあいまいさは発生しません。
「Collect Indices」配列（最後の例では X、Y）から抽出されたスライスの要素が入力配列ランクよりも少ない場合があり、明示的なマッピングにより、入力と同じランクになるようにインデックスを拡張する方法が指示されます。

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

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

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] 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
scatter_dims_to_operand_dims マップを使用して S を分散し、S を使用して operand にインデックス S_in を作成します。正式名称は次のとおりです。
1. Sk < scatter_dims_to_operand_dims.size の場合、_in[scatter_dims_to_operand_dims[k]] = S[k]。
2. Sそれ以外は _in[_] = 0。
inserted_window_dims に従って U の update_window_dims にインデックスを分散し、各 operands 配列にインデックス W_in を作成します。正式名称は次のとおりです。
1. 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} です）。
2. 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 が可換でない場合に重要です。

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` をそのまま使用するか、転置するか、共役の転置を取るかを指定します。

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

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

AfterAll

AllGather

AllReduce

AllToAll

BatchNormGrad

BatchNormInference

BatchNormTraining

BitcastConvertType

幅の異なるプリミティブ型へのビットキャスト変換

配信

BroadcastInDim

電話

コレスキー

クランプ

閉じる

CollectivePermute

Concatenate

条件

コンバージョン（畳み込み）

ConvWithGeneralPadding（畳み込み）

ConvertElementType

CrossReplicaSum

CustomCall

Dot

DotGeneral

DynamicSlice

DynamicUpdateSlice

要素ごとの 2 項算術演算

要素単位の比較演算

要素単位の単項関数

FFT

多次元 FFT

実装の詳細

収集

一般的なセマンティクス

非公式の説明と例

GetDimensionSize

SetDimensionSize

GetTupleElement

インフィード

ロッタ

マッピング

OptimizationBarrier

パッド

受信

削減

例

Variadic Reduce

ReducePrecision

ReduceScatter

ReduceWindow

ReplicaId

Reshape

巻き戻し（反転）

RngNormal

RngUniform

RngBitGenerator

散布

選択

SelectAndScatter

送信

スライス

並べ替え

行 / 列の入れ替え

TriangularSolve

タプル

一方

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