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

以下では、XlaBuilder インターフェースで定義されたオペレーションのセマンティクスについて説明します。通常、これらのオペレーションは、xla_data.proto の RPC インターフェースで定義されたオペレーションに 1 対 1 で対応します。

命名法に関する注意事項: XLA で一般化されたデータ型が扱うのは、32 ビット浮動小数点数など、均一な型の要素を保持する N 次元配列です。このドキュメントでは、配列を使用して任意の次元の配列を表します。便宜上、特殊なケースにはより具体的でわかりやすい名前が付けられています。たとえば、ベクトルは 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 は空にするか(この場合、すべてのレプリカは単一のグループに属します)、またはレプリカの数と同じ数の要素を含める必要があります。たとえば、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 ループがレプリカ間で何度も繰り返される場合があります。

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

この例では、4 つのコアが Alltoall に参加しています。各コアで、オペランドはディメンション 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 イプシロン値
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 複素数または浮動小数点型の、ランク > 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 つの配列はすべて同じ形状である必要があります。または、ブロードキャストの制限された形式として、min または maxT 型のスカラーにできます。

スカラー 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 のディメンションの連続するサブセットを並べます。

fold は、オペランドのディメンションの指定されたサブセットを単一のディメンションに置き換えます。入力引数は、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(operands..., dimension)

引数 タイプ セマンティクス
operands N XlaOp のシーケンス T 型の次元が [L0, L1, ...] の配列。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\) 型の単一の引数を取る必要があり、同じ型である false_operand で呼び出されます。true_computationfalse_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_indexS32 で、< 0 または >= N の場合、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 として指定されます。同じパディングでは、入力(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 ディメンションのサイズ(L 単位)と等しくなります。
  • 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 値が入っています。これは、畳み込みの場合はゼロを意味します。

右辺の拡張は、アトロス コンボリューションとも呼ばれます。詳しくは、tf.nn.atrous_conv2d をご覧ください。lh の拡散は転置畳み込みとも呼ばれます。詳しくは、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 ディメンションのサイズ(L 単位)と等しくなります。
  • z: カーネル(rhs)の output-z と同じサイズ。
  • spatial_dims: 畳み込みウィンドウの有効なプレースメントごとに 1 つの値。

上の図は、batch_group_count フィールドの仕組みを示しています。実質的に、各 1 分のバッチを batch_group_count グループにスライスし、出力特徴についても同じ処理をします。次に、これらのグループごとにペアワイズ畳み込みを行い、出力特徴ディメンションに沿って出力を連結します。他のすべてのディメンション(特徴と空間)のオペレーション セマンティクスは、すべて同じです。

畳み込みウィンドウの有効な配置は、ストライドと、パディング後のベース領域のサイズによって決まります。

畳み込み処理の動作を説明するには、2 次元の畳み込みについて考え、出力で batchzyx の固定座標を選択します。(y,x) は、ベース領域内のウィンドウの角の位置です(空間次元の解釈方法によっては左上隅など)。これで、ベース領域から取得された 2d ウィンドウが作成され、各 2d 点が 1d ベクトルに関連付けられているため、3d ボックスになります。畳み込みカーネルでは、出力座標 z を固定したため、3D ボックスもあります。2 つのボックスの寸法は同じであるため、2 つのボックス間の要素単位の積の合計を取得できます(ドット積と同様)。これが出力値です。

なお、output-z が例の場合、5 に設定すると、ウィンドウの各位置から、出力の 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 から浮動小数点への変換ルーチン(例: 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 の最初のディメンションに対する積の合計を実行します。これらは「契約」ディメンションです。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] バッチ MaMul
[b0, b1, m, k] dot [b0, b1, k, n] [b0, b1, m, n] バッチ MaMul

結果として、結果のディメンション番号は、バッチ ディメンションから始まり、次に 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} }

要素ごとのバイナリ算術演算

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)

ここで、Op は、Eq(等しい)、Ne(等しくない)、Ge(大なりイコール)、Gt(大なりイコール)、Le(不等号)、Lt(小なり)のいずれかです。EqTotalOrder、NeTotalOrder、GeTotalOrder、GtTotalOrder、LeTotalOrder、LtTotalOrder の演算子セットは、同じ機能を提供しますが、-NaN < -Inf < -Finite < -0 < +0 < +Finite < +Finite < N

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

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

次のオペレーションに対しては、異なるランクのブロードキャストをサポートする代替バリアントが存在します。

Op(lhs, rhs, broadcast_dimensions)

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

追加の broadcast_dimensions オペランドは、オペランドのブロードキャストに使用するディメンションを指定する整数のスライスです。セマンティクスの詳細については、ブロードキャストのページをご覧ください。

要素単位の単項関数

XlaBuilder は、次の要素単位の単項関数をサポートしています。

Abs(operand) 要素ごとの絶対値 x -> |x|

Ceil(operand) 要素単位のセル 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)にできます。

フィート

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 インデックスが呼び出し元によって確実に並べ替えられるかどうか。
unique_indices 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_sizes[k] を選択し、adjusted_slice_sizesslice_sizes であり、インデックス collapsed_slice_dims の境界は除去されます)。

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

  1. G = { Out[k] for batch_dims in 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 を使用して開始インデックス Sinoperand に作成します。より正確に言うと、

    1. Sin[start_index_map[k]] = S[k]: k < start_index_map.size の場合。

    2. Sin[_] = 0

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

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

    2. Oin[_] = 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 の昇順で)並べ替えられていると想定できます。含まれていない場合、セマンティクスは実装が定義されています。

unique_indices が true に設定されている場合、XLA は分散したすべての要素が一意であると想定できます。そのため、XLA では非アトミック演算を使用できます。unique_indices が true に設定され、分散されているインデックスが一意でない場合、セマンティクスは実装が定義されています。

非公式の説明と例

非公式には、出力配列のすべてのインデックス Out は、次のように計算されるオペランド配列の要素 E に対応します。

  • Out のバッチ ディメンションを使用して、start_indices から開始インデックスを検索します。

  • start_index_map を使用して、開始インデックス(operand.rank より小さい場合もある)を operand の「full」開始インデックスにマッピングします。

  • 完全な開始インデックスを使用して、サイズ slice_sizes のスライスを動的にスライスします。

  • スライスの形状を変更するには、collapsed_slice_dims ディメンションを折りたたみます。折りたたみスライスのサイズはすべて、境界 1 でなければならないため、この形状変更は常に有効です。

  • Out のオフセット ディメンションを使用してこのスライスにインデックスを付け、出力インデックス Out に対応する入力要素 E を取得します。

以下のすべての例で、index_vector_dimstart_indices.rank1)に設定されています。index_vector_dim に値を設定しても演算は根本的に変わりませんが、視覚的表現が煩雑になります。

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

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

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

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

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

収集インデックスは多次元である可能性があります。たとえば、上記の例のより一般的なバージョンで形状 [4,5,2] の「gather indices」配列を使用すると、次のようなインデックスが変換されます。

ここでも、バッチ動的スライスの G0 と、バッチ ディメンションの G1 として機能します。スライスのサイズは [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 つだけあります。同様に、値が O0 である出力オフセット インデックスは 1 つだけです。ただし、入力配列のインデックスとして使用される前に、これらは「一連のインデックス マッピング」(正式な説明の start_index_map)と「オフセット マッピング」(正式な説明の remapped_offset_dims)に従って展開され、それぞれ [X,0] と [0O0] に展開され、[XG1GOG1GOG0O0]01 にセマンティック インデックス [0G1GOG0O0]01 にセマンティックGatherIndicestf.gather_nd

このケースの slice_sizes[1,11] です。これは直感的に、収集インデックス配列のすべてのインデックス 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 つの計算で複数のインフィード処理を使用できますが、インフィード処理の順序は 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 ]]

Map

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

Map(operands..., computation)

引数 タイプ セマンティクス
operands N 個の XlaOp のシーケンス T0..T{N-1} 型の N 個の配列
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_padding 値と interior_padding 値の例を示しています。

Recv

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 命令によって作成されたコンテキストが与えられると、データ転送が完了するまで待機し、受信したデータを返します。

Reduce(減らす)

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、ディメンション 0 はサイズ 2、ディメンション 1 はサイズ 3 です。

「add」関数でディメンション 0 または 1 を縮小した結果:

なお、リダクション結果はどちらも 1 次元配列です。この図では、見やすくするため、一方を列、もう一方を行として示しています。

より複雑な例として、3D 配列を示します。ランクは 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 の場合、すべての入力に同時に適用されるため、Reduce 関数の適用はやや複雑になります。オペランドは、次の順序で計算に提供されます。

  • 第 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 ブロックに分割し、レプリカ グループ内のレプリカ iith シャードを受け取る集合オペレーションです。

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 分の 1 に縮小された入力シェイプです。たとえば、2 つのレプリカがあり、2 つのレプリカでオペランドの値が [1.0, 2.25][3.0, 5.25] である場合、scatter_dim0 であるこの演算からの出力値は、最初のレプリカでは [4.0]、2 番目のレプリカでは [7.5] になります。

ReduceWindow

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

N 個の多次元配列のシーケンスの各ウィンドウ内のすべての要素にリダクション関数を適用し、出力として 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 ウィンドウのパディング タイプ(ストライドが 1 の場合に入力と同じ出力形状になるようにパディングする Padding::kSame。またはパディングを使用せず、ウィンドウが収まらないと「停止」する 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 を使用してレデューサ ウィンドウの最小値(初期値は MAX_FLOAT)を計算し、入力配列 [10000, 1000, 100, 10, 1] でストライド 2 を計算することを検討してください。パディング kValid により、2 つの有効なウィンドウ([10000, 1000, 100][100, 10, 1])の最小値が計算され、出力は [100, 1] になります。まず kSame をパディングすると、両側に初期要素を追加して [MAX_VALUE, 10000, 1000, 100, 10, 1, MAX_VALUE] を得ることで、Reduce-window の後の形状がストライド 1 の入力と同じになるように配列がパディングされます。パディングされた配列に対して Reduce-window を実行すると、3 つのウィンドウ [MAX_VALUE, 10000, 1000][1000, 100, 10][10, 1, MAX_VALUE] で動作し、[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(f32[1x1] { {5} }, {0,1}, {}) == 5;
Reshape(5, {}, {1,1}) == f32[1x1] { {5} };

Rev(リバース)

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 の値である一連の結果を生成します。update_computation を使用して、updates の値のシーケンスで更新される複数のスライス(scatter_indices で指定されたインデックス)があります。

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 インデックスが呼び出し元によって確実に並べ替えられるかどうか。

ここで

  • 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 に存在する場合(つまり、一部の kupdate_window_dims[k] と等しい場合)、updates のディメンション i の境界は、inserted_window_dims を考慮した後、operand の対応する境界を超えてはなりません(adjusted_window_bounds[k]。adjusted_window_bounds[k] には、インデックス inserted_window_dims の境界が削除された operand の境界が含まれます)。
    • iupdate_scatter_dims に存在する場合(つまり、一部の kupdate_scatter_dims [k] と等しい場合)、updates のディメンション i の境界は、index_vector_dim をスキップして、対応する scatter_indices の境界と等しくなければなりません(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) は位置 index_vector_dim の b を A に挿入します。k
  2. scatter_dims_to_operand_dims のマップを使用して S を分散し、S を使用して operand にインデックス Sin を作成します。より正式に次のように記述します。
    1. Sin[scatter_dims_to_operand_dims[k]] = S[k](k < scatter_dims_to_operand_dims.size の場合)。
    2. Sin[_] = 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. Win[_] = 0
  4. IWin + Sin です。「+」は要素単位の加算です。

要約すると、scatter 操作は次のように定義できます。

  • outputoperands で初期化します。つまり、operands[J] 配列内のすべてのインデックス O について、すべてのインデックス J に対して初期化します。
    output[J][O] = operands[J][O]
  • updates[J] 配列のすべてのインデックス U と、operand[J] 配列の対応するインデックス OOoutput の有効なインデックスの場合):
    (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 の昇順で)並べ替えられていると想定できます。含まれていない場合、セマンティクスは実装が定義されています。

非公式には、scatter 演算は collect 演算の逆とみなすことができます。すなわち、scatter 演算は、対応する collect 演算によって抽出された入力内の要素を更新します。

詳細な非公式の説明と例については、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, 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 オペランドは、同じ次元のテンソルでなければなりません。テンソルの要素型が異なる場合があります。

  • すべてのオペランドは、個別にではなく、まとめて並べ替えられます。概念的には、オペランドはタプルとして扱われます。インデックス位置 ij にある各オペランドの要素を入れ替える必要があるかどうかをチェックする場合は、2 * n スカラー パラメータでコンパレータが呼び出されます。ここで、パラメータ 2 * kk-th オペランドの i 位置の値に対応し、パラメータ 2 * k + 1k-th オペランドの位置 j の値に対応します。したがって、通常、コンパレータはパラメータ 2 * k2 * k + 1 を相互に比較し、他のパラメータペアをタイブレーカーとして使用する可能性があります。

  • 結果は、(上記のように指定されたディメンションに沿って)並べ替えられたオペランドで構成されるタプルになります。タプルの i-th オペランドは、Sort の i-th オペランドに対応します。

たとえば、operand0 = [3, 1]operand1 = [42, 50]operand2 = [-3.0, 1.1] の 3 つのオペランドがあり、コンパレータが 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, permut, Permute(permut, operand.shape.dimensions)) と同じです。

TriangularSolve

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

前方代入または後代入法によって、下または上三角係数行列を持つ一次方程式系を解く。リーディング ディメンションに沿ってブロードキャストし、このルーティンは、ab が与えられた変数 x に対して、行列系 op(a) * x = b または x * op(a) = b のいずれかを解きます。ここで、op(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 型は Tuple で、反復回数の int32 とアキュムレータの vector[10] で構成されています。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};
}