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

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

命名法に関する注: XLA が扱う一般化されたデータ型は、均一な型（32 ビット浮動小数点など）の要素を保持する N 次元配列です。このドキュメントでは、任意の次元の配列を示すために array が使用されています。便宜上、特殊なケースにはより具体的でわかりやすい名前が付けられています。たとえば、ベクトルは 1 次元配列、行列は 2 次元配列です。

Op の構造について詳しくは、シェイプとレイアウトとタイル レイアウトをご覧ください。

腹筋

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

要素ごとの絶対値 x -> |x|。

Abs(operand)

StableHLO の詳細については、StableHLO - abs をご覧ください。

追加

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

lhs と rhs の要素ごとの加算を行います。

Add(lhs, rhs)

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

Add には、異なる次元のブロードキャストをサポートする代替バリアントがあります。

Add(lhs,rhs, broadcast_dimensions)

このオペレーションのバリアントは、ランクの異なる配列間の算術演算（行列とベクトルの加算など）に使用する必要があります。

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

StableHLO の情報については、StableHLO - add をご覧ください。

AddDependency

HloInstruction::AddDependency もご覧ください。

AddDependency は HLO ダンプに表示されることがありますが、エンドユーザーが手動で構築することを意図したものではありません。

AfterAll

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

AfterAll は可変数のトークンを受け取り、単一のトークンを生成します。トークンは、副作用のあるオペレーション間でスレッド化して順序を強制できるプリミティブ型です。AfterAll は、一連のオペレーションの後にオペレーションを順序付けるためのトークンの結合として使用できます。

AfterAll(tokens)

StableHLO の情報については、StableHLO - after_all をご覧ください。

AllGather

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

レプリカ間で連結を実行します。

AllGather(operand, all_gather_dimension, shard_count, replica_groups, channel_id, layout, use_global_device_ids)

• replica_groups は、連結が実行されるレプリカ グループのリストです（現在のレプリカのレプリカ ID は ReplicaId を使用して取得できます）。各グループ内のレプリカの順序によって、結果内の入力の順序が決まります。replica_groups は空（この場合、すべてのレプリカは単一のグループに属し、0 から N - 1 の順に並べられます）にするか、レプリカの数と同じ数の要素を含める必要があります。たとえば、replica_groups = {0, 2}, {1, 3} はレプリカ 0 と 2、1 と 3 の間で連結を実行します。
• shard_count は各レプリカ グループのサイズです。これは、replica_groups が空の場合に必要です。
• channel_id はモジュール間の通信に使用されます。同じ channel_id を持つ all-gather オペレーションのみが相互に通信できます。
• use_global_device_ids ReplicaGroup 構成の ID がレプリカ ID ではなく、（replica_id * partition_count + partition_id）のグローバル ID を表す場合は true を返します。これにより、このオールリデュースがクロス パーティションとクロス レプリカの両方である場合、デバイスのグループ化をより柔軟に行うことができます。

出力の形状は、入力の形状で all_gather_dimension が shard_count 倍に拡大されたものです。たとえば、2 つのレプリカがあり、オペランドの値が 2 つのレプリカでそれぞれ [1.0, 2.5] と [3.0, 5.25] の場合、all_gather_dim が 0 のこの op からの出力値は、両方のレプリカで [1.0, 2.5, 3.0,5.25] になります。

AllGather の API は、内部的に 2 つの HLO 命令（AllGatherStart と AllGatherDone）に分解されます。

HloInstruction::CreateAllGatherStart もご覧ください。

AllGatherStart、AllGatherDone は HLO のプリミティブとして機能します。これらのオペレーションは HLO ダンプに表示されることがありますが、エンドユーザーが手動で構築することを想定していません。

StableHLO の詳細については、StableHLO - all_gather をご覧ください。

AllReduce

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

レプリカ間でカスタム計算を実行します。

AllReduce(operand, computation, replica_groups, channel_id, shape_with_layout, use_global_device_ids)

• operand が配列のタプルの場合、タプルの各要素に対して all-reduce が実行されます。
• replica_groups は、削減が実行されるレプリカ グループのリストです（現在のレプリカのレプリカ ID は ReplicaId を使用して取得できます）。replica_groups は空（すべてのレプリカが単一のグループに属している場合）であるか、レプリカの数と同じ数の要素を含んでいる必要があります。たとえば、replica_groups = {0, 2}, {1, 3} はレプリカ 0 と 2、1 と 3 の間で削減を実行します。
• channel_id はモジュール間の通信に使用されます。同じ channel_id を持つ all-reduce オペレーションのみが相互に通信できます。
• shape_with_layout: AllReduce のレイアウトを指定されたレイアウトに強制します。これは、個別にコンパイルされた AllReduce オペレーションのグループに同じレイアウトを保証するために使用されます。
• use_global_device_ids ReplicaGroup 構成の ID がレプリカ ID ではなく、（replica_id * partition_count + partition_id）のグローバル ID を表す場合は true を返します。これにより、このオールリデュースがクロス パーティションとクロス レプリカの両方である場合、デバイスのグループ化をより柔軟に行うことができます。

出力の形状は入力の形状と同じです。たとえば、2 つのレプリカがあり、オペランドの値が 2 つのレプリカでそれぞれ [1.0, 2.5] と [3.0, 5.25] の場合、この op と合計計算の出力値は両方のレプリカで [4.0, 7.75] になります。入力がタプルの場合、出力もタプルになります。

AllReduce の結果を計算するには、各レプリカからの入力が 1 つ必要です。そのため、あるレプリカが別のレプリカよりも AllReduce ノードを多く実行すると、前者のレプリカは永久に待機します。レプリカはすべて同じプログラムを実行しているため、このような状況が発生する可能性はそれほど高くありませんが、while ループの条件が infeed のデータに依存し、その infeed のデータによって、あるレプリカで while ループが別のレプリカよりも多く繰り返される場合に発生する可能性があります。

AllReduce の API は、内部的に 2 つの HLO 命令（AllReduceStart と AllReduceDone）に分解されます。

HloInstruction::CreateAllReduceStart もご覧ください。

AllReduceStart と AllReduceDone は HLO のプリミティブとして機能します。これらのオペレーションは HLO ダンプに表示されることがありますが、エンドユーザーが手動で構築することを想定していません。

CrossReplicaSum

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

合計計算で AllReduce を実行します。

CrossReplicaSum(operand, replica_groups)

レプリカの各サブグループ内のオペランド値の合計を返します。すべてのレプリカが合計に 1 つの入力を提供し、すべてのレプリカが各サブグループの結果の合計を受け取ります。

AllToAll

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

AllToAll は、すべてのコアからすべてのコアにデータを送信する集団オペレーションです。次の 2 つのフェーズがあります。

1. 散開フェーズ。各コアで、オペランドは split_dimensions に沿って split_count 個のブロックに分割され、ブロックはすべてのコアに分散されます。たとえば、i 番目のブロックは i 番目のコアに送信されます。
2. 収集フェーズ。各コアは、受信したブロックを concat_dimension に沿って連結します。

参加するコアは、次の方法で構成できます。

• replica_groups: 各 ReplicaGroup には、計算に参加しているレプリカ ID のリストが含まれています（現在のレプリカのレプリカ ID は ReplicaId を使用して取得できます）。AllToAll は、指定された順序でサブグループ内で適用されます。たとえば、replica_groups = { {1,2,3}, {4,5,0} } は、レプリカ {1, 2, 3} 内の gather フェーズで AllToAll が適用され、受信したブロックが 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, layout, channel_id)

シェイプとレイアウトの詳細については、xla::shapes をご覧ください。

StableHLO の詳細については、StableHLO - all_to_all をご覧ください。

AllToAll - 例 1。

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 つあります。各コアで、オペランドはディメンション 1 に沿って 4 つの部分に分割されるため、各部分の形状は f32[4,4] になります。4 つの部分はすべてのコアに分散されます。次に、各コアは、コア 0 ～ 4 の順に、受け取った部分をディメンション 0 に沿って連結します。したがって、各コアの出力の形状は f32[16,4] になります。

AllToAll - 例 2 - StableHLO

上記の例では、AllToAll に 2 つのレプリカが参加しています。各レプリカで、オペランドの形状は f32[2,4] です。オペランドはディメンション 1 に沿って 2 つの部分に分割されるため、各部分の形状は f32[2,2] になります。2 つの部分は、レプリカ グループ内の位置に従ってレプリカ間で交換されます。各レプリカは、両方のオペランドから対応する部分を収集し、ディメンション 0 に沿って連結します。その結果、各レプリカの出力の形状は f32[4,2] になります。

RaggedAllToAll

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

RaggedAllToAll は、入力と出力がラグド テンソルである集団 all-to-all オペレーションを実行します。

RaggedAllToAll(input, input_offsets, send_sizes, output, output_offsets, recv_sizes, replica_groups, channel_id)

ラグド テンソルは、次の 3 つのテンソルのセットで定義されます。

• data: data テンソルは最も外側のディメンションに沿って「ラグ」があり、各インデックス付き要素のサイズは可変です。
• offsets: offsets テンソルは data テンソルの最も外側のディメンションをインデックス化し、data テンソルの各ラグド要素の開始オフセットを表します。
• sizes: sizes テンソルは、data テンソルの各ラグド要素のサイズを表します。サイズはサブ要素の単位で指定されます。サブ要素は、最も外側の「ラグ」次元を削除して取得した「データ」テンソル形状の接尾辞として定義されます。
• offsets テンソルと sizes テンソルは同じサイズにする必要があります。

ラグド テンソルの例:

data: [8,3] =
{ {a,b,c},{d,e,f},{g,h,i},{j,k,l},{m,n,o},{p,q,r},{s,t,u},{v,w,x} }

offsets: [3] = {0, 1, 4}

sizes: [3] = {1, 3, 4}

// Index 'data' at 'offsets'[0], 'sizes'[0]' // {a,b,c}

// Index 'data' at 'offsets'[1], 'sizes'[1]' // {d,e,f},{g,h,i},{j,k,l}

// Index 'data' at 'offsets'[2], 'sizes'[2]' // {m,n,o},{p,q,r},{s,t,u},{v,w,x}

output_offsets は、各レプリカがターゲット レプリカ出力パースペクティブでオフセットを持つようにシャーディングする必要があります。

i 番目の出力オフセットの場合、現在のレプリカは input[input_offsets[i]:input_offsets[i]+input_sizes[i]] 更新を i 番目のレプリカに送信します。これは、i 番目のレプリカの output の output_i[output_offsets[i]:output_offsets[i]+send_sizes[i]] に書き込まれます。

たとえば、レプリカが 2 つある場合:

replica 0:
input: [1, 2, 2]
output:[0, 0, 0, 0]
input_offsets: [0, 1]
send_sizes: [1, 2]
output_offsets: [0, 0]
recv_sizes: [1, 1]

replica 1:
input: [3, 4, 0]
output: [0, 0, 0, 0]
input_offsets: [0, 1]
send_sizes: [1, 1]
output_offsets: [1, 2]
recv_sizes: [2, 1]

// replica 0's result will be: [1, 3, 0, 0]
// replica 1's result will be: [2, 2, 4, 0]

ラグド all-to-all HLO には次の引数があります。

• input: 不規則な入力データ テンソル。
• output: ラグド出力データ テンソル。
• input_offsets: 不規則な入力オフセット テンソル。
• send_sizes: 不規則な送信サイズ テンソル。
• output_offsets: ターゲット レプリカ出力の不均一なオフセットの配列。
• recv_sizes: ラグド recv サイズ テンソル。

*_offsets テンソルと *_sizes テンソルはすべて同じ形状である必要があります。

*_offsets テンソルと *_sizes テンソルでは、次の 2 つのシェイプがサポートされています。

• [num_devices]。ragged-all-to-all は、レプリカ グループ内の各リモート デバイスに最大 1 つの更新を送信できます。次に例を示します。

for (remote_device_id : replica_group) {
     SEND input[input_offsets[remote_device_id]],
     output[output_offsets[remote_device_id]],
     send_sizes[remote_device_id] }

• [num_devices, num_updates] では、レプリカ グループ内の各リモート デバイスに対して、最大 num_updates 個の更新を同じリモート デバイスに送信できます（それぞれ異なるオフセット）。

次に例を示します。

for (remote_device_id : replica_group) {
    for (update_idx : num_updates) {
        SEND input[input_offsets[remote_device_id][update_idx]],
        output[output_offsets[remote_device_id][update_idx]]],
        send_sizes[remote_device_id][update_idx] } }

および

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

2 つのテンソル lhs と rhs の要素ごとの AND を実行します。

And(lhs, rhs)

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

And には、異なる次元のブロードキャストをサポートする別のバリアントがあります。

And(lhs,rhs, broadcast_dimensions)

このオペレーションのバリアントは、ランクの異なる配列間の算術演算（行列とベクトルの加算など）に使用する必要があります。

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

StableHLO の詳細については、StableHLO - and をご覧ください。

非同期

HloInstruction::CreateAsyncStart、HloInstruction::CreateAsyncUpdate、HloInstruction::CreateAsyncDone もご覧ください。

AsyncDone、AsyncStart、AsyncUpdate は、非同期オペレーションに使用される内部 HLO 命令であり、HLO のプリミティブとして機能します。これらのオペレーションは HLO ダンプに表示されることがありますが、エンドユーザーが手動で構築することを想定していません。

Atan2

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

lhs と rhs で要素ごとの atan2 演算を実行します。

Atan2(lhs, rhs)

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

Atan2 には、異なる次元のブロードキャストをサポートする代替バリアントがあります。

Atan2(lhs,rhs, broadcast_dimensions)

このオペレーションのバリアントは、ランクの異なる配列間の算術演算（行列とベクトルの加算など）に使用する必要があります。

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

StableHLO の詳細については、StableHLO - atan2 をご覧ください。

BatchNormGrad

アルゴリズムの詳細については、XlaBuilder::BatchNormGrad とバッチ正規化に関する元の論文もご覧ください。

バッチ正規化の勾配を計算します。

BatchNormGrad(operand, scale, batch_mean, batch_var, grad_output, epsilon, feature_index)

特徴ディメンションの各特徴（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} \]

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

出力タイプは、次の 3 つのハンドルのタプルです。

StableHLO の情報については、StableHLO - batch_norm_grad をご覧ください。

BatchNormInference

アルゴリズムの詳細については、XlaBuilder::BatchNormInference とバッチ正規化に関する元の論文もご覧ください。

バッチ ディメンションと空間ディメンションにわたって配列を正規化します。

BatchNormInference(operand, scale, offset, mean, variance, epsilon, feature_index)

特徴量ディメンションの各特徴量について（feature_index は operand の特徴量ディメンションのインデックス）、このオペレーションは他のすべてのディメンションの平均と分散を計算し、その平均と分散を使用して operand の各要素を正規化します。feature_index は、operand の特徴ディメンションの有効なインデックスである必要があります。

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

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

StableHLO の詳細については、StableHLO - batch_norm_inference をご覧ください。

BatchNormTraining

アルゴリズムの詳細については、XlaBuilder::BatchNormTraining と the original batch normalization paper もご覧ください。

バッチ ディメンションと空間ディメンションにわたって配列を正規化します。

BatchNormTraining(operand, scale, offset, epsilon, feature_index)

特徴量ディメンションの各特徴量について（feature_index は operand の特徴量ディメンションのインデックス）、このオペレーションは他のすべてのディメンションの平均と分散を計算し、その平均と分散を使用して operand の各要素を正規化します。feature_index は、operand の特徴ディメンションの有効なインデックスである必要があります。

w と h を空間次元のサイズとして含む m 要素を含む operand \(x\) の各バッチのアルゴリズムは次のとおりです（operand は 4 次元配列と仮定）。

• 特徴ディメンションの各特徴 l のバッチ平均 \(\mu_l\) を計算します。 \(\mu_l=\frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h x_{ijkl}\)

• バッチ分散を計算します。 \(\sigma^2_l\): $\sigma^2l=\frac{1}{mwh}\sum{i=1}^m\sum{j=1}^w\sum{k=1}^h (x_{ijkl} - \mu_l)^2$

• 正規化、スケーリング、シフト:\(y_{ijkl}=\frac{\gamma_l(x_{ijkl}-\mu_l)}{\sqrt[2]{\sigma^2_l+\epsilon} }+\beta_l\)

通常は小さい数値であるイプシロン値が追加され、ゼロ除算エラーが回避されます。

出力の型は、3 つの XlaOp のタプルです。

batch_mean と batch_var は、上記の式を使用してバッチ ディメンションと空間ディメンション全体で計算されたモーメントです。

StableHLO の情報については、StableHLO - batch_norm_training をご覧ください。

Bitcast

HloInstruction::CreateBitcast もご覧ください。

Bitcast は HLO ダンプに表示されることがありますが、エンドユーザーが手動で構築することを意図したものではありません。

BitcastConvertType

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

TensorFlow の tf.bitcast と同様に、データシェイプからターゲット シェイプへの要素ごとのビットキャスト オペレーションを実行します。入力と出力のサイズは一致している必要があります。たとえば、s32 要素は bitcast ルーチンを介して f32 要素になり、1 つの s32 要素は 4 つの s8 要素になります。Bitcast は低レベルのキャストとして実装されるため、浮動小数点表現が異なるマシンでは異なる結果が得られます。

BitcastConvertType(operand, new_element_type)

オペランドとターゲット シェイプのディメンションは、変換前後のプリミティブ サイズの比率で変化する最後のディメンションを除き、一致している必要があります。

移行元と移行先の要素の型はタプルにできません。

StableHLO の詳細については、StableHLO - bitcast_convert をご覧ください。

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

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)

新しいディメンションは左側に挿入されます。つまり、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 になります。

StableHLO の詳細については、StableHLO - ブロードキャストをご覧ください。

BroadcastInDim

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

配列内のデータを複製して、配列のサイズとディメンションの数を拡大します。

BroadcastInDim(operand, out_dim_size, broadcast_dimensions)

ブロードキャストに似ていますが、任意の場所にディメンションを追加したり、サイズ 1 の既存のディメンションを拡張したりできます。

operand は、out_dim_size で記述された形状にブロードキャストされます。broadcast_dimensions は、operand のディメンションをターゲット形状のディメンションにマッピングします。つまり、オペランドの i 番目のディメンションは、出力形状の broadcast_dimension[i] 番目のディメンションにマッピングされます。operand のディメンションのサイズは 1 であるか、マッピング先の出力シェイプのディメンションと同じサイズである必要があります。残りのディメンションはサイズ 1 のディメンションで埋められます。縮退次元ブロードキャストは、これらの縮退次元に沿ってブロードキャストし、出力形状に到達します。セマンティクスについては、ブロードキャスト ページで詳しく説明しています。

電話

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

指定された引数を使用して計算を呼び出します。

Call(computation, operands...)

operands の項数と型は、computation のパラメータと一致する必要があります。operands がなくても構いません。

CompositeCall

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

他の StableHLO オペレーションで構成（合成）されたオペレーションをカプセル化し、入力と composite_attributes を受け取って結果を生成します。演算のセマンティクスは分解属性によって実装されます。複合演算は、プログラムのセマンティクスを変更することなく、その分解に置き換えることができます。分解をインライン化しても同じオペレーション セマンティクスが得られない場合は、custom_call の使用を優先します。

バージョン フィールド（デフォルトは 0）は、複合のセマンティクスが変更されたタイミングを示すために使用されます。

このオペレーションは、属性 is_composite=true を持つ kCall として実装されます。decomposition フィールドは computation 属性で指定します。フロントエンド属性には、composite. という接頭辞が付いた残りの属性が格納されます。

CompositeCall オペレーションの例:

f32[] call(f32[] %cst), to_apply=%computation, is_composite=true,
frontend_attributes = {
  composite.name="foo.bar",
  composite.attributes={n = 1 : i32, tensor = dense<1> : tensor<i32>},
  composite.version="1"
}

CompositeCall(computation, operands..., name, attributes, version)

op の decomposition は呼び出されるフィールドではなく、下位レベルの実装（to_apply=%funcname など）を含む関数を指す to_apply 属性として表示されます。

合成と分解の詳細については、StableHLO 仕様をご覧ください。

Cbrt

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

要素ごとの立方根演算 x -> cbrt(x)。

Cbrt(operand)

Cbrt は、オプションの result_accuracy 引数もサポートしています。

Cbrt(operand, result_accuracy)

result_accuracy の詳細については、結果の精度をご覧ください。

StableHLO の詳細については、StableHLO - cbrt をご覧ください。

切り上げ

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

要素ごとの天井関数 x -> ⌈x⌉。

Ceil(operand)

StableHLO の情報については、StableHLO - ceil をご覧ください。

Cholesky

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

対称（エルミート）正定値行列のバッチの Cholesky 分解を計算します。

Cholesky(a, lower)

lower が true の場合、$a = l$ となる下三角行列 l を計算します。l^T$。lower が false の場合、\(a = u^T . u\)となる上三角行列 u を計算します。

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

a の次元が 2 より大きい場合、a は行列のバッチとして扱われます。ここで、下位 2 次元を除くすべての次元がバッチ ディメンションになります。

a が対称（エルミート）正定値でない場合、結果は実装定義になります。

StableHLO の詳細については、StableHLO - cholesky をご覧ください。

範囲制限

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

オペランドを最小値と最大値の範囲内にクランプします。

Clamp(min, operand, max)

オペランドと最小値、最大値が指定された場合、オペランドが最小値と最大値の範囲内であればオペランドを返し、オペランドがこの範囲を下回る場合は最小値を返し、オペランドがこの範囲を上回る場合は最大値を返します。つまり、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};

StableHLO の詳細については、StableHLO - clamp をご覧ください。

閉じる

XlaBuilder::Collapse もご覧ください。tf.reshape オペレーション。

配列のディメンションを 1 つのディメンションに折りたたみます。

Collapse(operand, dimensions)

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

Clz

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

要素ごとに先頭のゼロの数をカウントします。

Clz(operand)

CollectiveBroadcast

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

レプリカ間でデータをブロードキャストします。各グループの最初のレプリカ ID から、同じグループ内の他の ID にデータが送信されます。レプリカ ID がレプリカ グループにない場合、そのレプリカの出力は shape の 0 で構成されるテンソルになります。

CollectiveBroadcast(operand, replica_groups, channel_id)

StableHLO の詳細については、StableHLO - collective_broadcast をご覧ください。

CollectivePermute

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

CollectivePermute は、レプリカ間でデータを送受信する集団オペレーションです。

CollectivePermute(operand, source_target_pairs, channel_id, inplace)

source_target_pairs には次の制限があります。

• 2 つのペアに同じターゲット レプリカ ID を使用することはできません。また、同じソース レプリカ ID を使用することもできません。
• レプリカ ID がペアのターゲットでない場合、そのレプリカの出力は、入力と同じ形状の 0 で構成されるテンソルになります。

CollectivePermute オペレーションの API は、内部的に 2 つの HLO 命令（CollectivePermuteStart と CollectivePermuteDone）に分解されます。

HloInstruction::CreateCollectivePermuteStart もご覧ください。

CollectivePermuteStart と CollectivePermuteDone は HLO のプリミティブとして機能します。これらのオペレーションは HLO ダンプに表示されることがありますが、エンドユーザーが手動で構築することを想定していません。

StableHLO の情報については、StableHLO - collective_permute をご覧ください。

比較

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

次の lhs と rhs の要素ごとの比較を行います。

Eq

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

lhs と rhs の要素ごとの等しい比較を行います。

\(lhs = rhs\)

Eq(lhs, rhs)

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

Eq には、異なる次元のブロードキャストをサポートする別のバリアントがあります。

Eq(lhs,rhs, broadcast_dimensions)

このオペレーションのバリアントは、ランクの異なる配列間の算術演算（行列とベクトルの加算など）に使用する必要があります。

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

Eq の浮動小数点数に対する全順序をサポートします。

\[-NaN < -Inf < -Finite < -0 < +0 < +Finite < +Inf < +NaN.\]

EqTotalOrder(lhs,rhs, broadcast_dimensions)

StableHLO の詳細については、StableHLO - 比較をご覧ください。

Ne

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

lhs と rhs の要素ごとの等しくない比較を行います。

\(lhs != rhs\)

Ne(lhs, rhs)

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

Ne には、異なる次元のブロードキャストをサポートする別のバリアントがあります。

Ne(lhs,rhs, broadcast_dimensions)

このオペレーションのバリアントは、ランクの異なる配列間の算術演算（行列とベクトルの加算など）に使用する必要があります。

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

Ne では、次の条件を適用することで、浮動小数点数を超える合計注文をサポートしています。

\[-NaN < -Inf < -Finite < -0 < +0 < +Finite < +Inf < +NaN.\]

NeTotalOrder(lhs,rhs, broadcast_dimensions)

StableHLO の詳細については、StableHLO - 比較をご覧ください。

Ge

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

lhs と rhs の要素ごとのgreater-or-equal-than比較を行います。

\(lhs >= rhs\)

Ge(lhs, rhs)

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

Ge には、異なる次元のブロードキャストをサポートする代替バリアントがあります。

Ge(lhs,rhs, broadcast_dimensions)

このオペレーションのバリアントは、ランクの異なる配列間の算術演算（行列とベクトルの加算など）に使用する必要があります。

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

Gt の浮動小数点数に対する全順序をサポートします。次の条件を適用します。

\[-NaN < -Inf < -Finite < -0 < +0 < +Finite < +Inf < +NaN.\]

GtTotalOrder(lhs,rhs, broadcast_dimensions)

StableHLO の詳細については、StableHLO - 比較をご覧ください。

Gt

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

lhs と rhs の要素ごとの大なり比較を行います。

\(lhs > rhs\)

Gt(lhs, rhs)

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

Gt には、異なる次元のブロードキャストをサポートする代替バリアントが存在します。

Gt(lhs,rhs, broadcast_dimensions)

このオペレーションのバリアントは、ランクの異なる配列間の算術演算（行列とベクトルの加算など）に使用する必要があります。

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

StableHLO の詳細については、StableHLO - 比較をご覧ください。

Le

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

lhs と rhs の要素ごとのless-or-equal-than比較を行います。

\(lhs <= rhs\)

Le(lhs, rhs)

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

Le には、異なる次元のブロードキャストをサポートする代替バリアントがあります。

Le(lhs,rhs, broadcast_dimensions)

このオペレーションのバリアントは、ランクの異なる配列間の算術演算（行列とベクトルの加算など）に使用する必要があります。

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

Le の浮動小数点数を超える合計注文をサポートします。次の条件を適用します。

\[-NaN < -Inf < -Finite < -0 < +0 < +Finite < +Inf < +NaN.\]

LeTotalOrder(lhs,rhs, broadcast_dimensions)

StableHLO の詳細については、StableHLO - 比較をご覧ください。

Lt

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

lhs と rhs の要素ごとの小なり比較を行います。

\(lhs < rhs\)

Lt(lhs, rhs)

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

Lt には、異なる次元のブロードキャストをサポートする代替バリアントがあります。

Lt(lhs,rhs, broadcast_dimensions)

このオペレーションのバリアントは、ランクの異なる配列間の算術演算（行列とベクトルの加算など）に使用する必要があります。

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

Lt の浮動小数点数を超える合計注文をサポートします。次の条件を適用します。

\[-NaN < -Inf < -Finite < -0 < +0 < +Finite < +Inf < +NaN.\]

LtTotalOrder(lhs,rhs, broadcast_dimensions)

StableHLO の詳細については、StableHLO - 比較をご覧ください。

複雑

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

実数値と虚数値のペア lhs と rhs から複素数値への要素ごとの変換を行います。

Complex(lhs, rhs)

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

Complex には、異なる次元のブロードキャストをサポートする代替バリアントがあります。

Complex(lhs,rhs, broadcast_dimensions)

このオペレーションのバリアントは、ランクの異なる配列間の算術演算（行列とベクトルの加算など）に使用する必要があります。

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

StableHLO の詳細については、StableHLO - 複素数をご覧ください。

ConcatInDim（連結）

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

Concatenate は、複数の配列オペランドから配列を作成します。配列の次元数は、入力配列オペランドの次元数と同じです（入力配列オペランドの次元数は互いに同じである必要があります）。配列には、指定された順序で引数が含まれます。

Concatenate(operands..., dimension)

dimension を除くすべてのディメンションは同じである必要があります。これは、XLA が「不揃い」な配列をサポートしていないためです。また、0 次元値は連結できません（連結が発生するディメンションに名前を付けることができないため）。

1 次元の例:

Concat({ {2, 3}, {4, 5}, {6, 7} }, 0)
//Output:  {2, 3, 4, 5, 6, 7}

2 次元の例:

let a = { {1, 2},
         {3, 4},
         {5, 6} };

let b = { {7, 8} };

Concat({a, b}, 0)

//Output:  { {1, 2},
//          {3, 4},
//          {5, 6},
//          {7, 8} }

図:

StableHLO の詳細については、StableHLO - 連結をご覧ください。

条件文

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

Conditional(predicate, true_operand, true_computation, false_operand, false_computation)

predicate が true の場合は true_computation を、predicate が false の場合は false_computation を実行し、結果を返します。

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

predicate の値に応じて、true_computation と false_computation のいずれか 1 つのみが実行されます。

Conditional(branch_index, branch_computations, branch_operands)

branch_computations[branch_index] を実行し、結果を返します。branch_index が 0 未満または N 以上の S32 の場合、branch_computations[N-1] はデフォルトのブランチとして実行されます。

各 branch_computations[b] は \(T_b\) 型の単一の引数を取り、同じ型の branch_operands[b] で呼び出されます。各 branch_computations[b] の戻り値の型は同じである必要があります。

branch_index の値に応じて、branch_computations のいずれか 1 つのみが実行されます。

StableHLO の情報については、StableHLO - if をご覧ください。

定数

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

定数 literal から output を生成します。

Constant(literal)

StableHLO の情報については、StableHLO - 定数をご覧ください。

ConvertElementType

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

C++ の要素ごとの static_cast と同様に、ConvertElementType はデータシェイプからターゲット シェイプへの要素ごとの変換オペレーションを実行します。ディメンションは一致している必要があり、変換は要素ごとに行われます。たとえば、s32 要素は s32 から f32 への変換ルーチンを介して f32 要素になります。

ConvertElementType(operand, new_element_type)

オペランドのディメンションとターゲット シェイプは一致している必要があります。ソース要素と宛先要素の型はタプルであってはなりません。

T=s32 から U=f32 への変換では、最近接偶数丸めなどの正規化 int から float への変換ルーチンが実行されます。

let a: s32[3] = {0, 1, 2};
let b: f32[3] = convert(a, f32);
then b == f32[3]{0.0, 1.0, 2.0}

StableHLO の詳細については、StableHLO - 変換をご覧ください。

Conv（畳み込み）

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

ニューラル ネットワークで使用される種類の畳み込みを計算します。ここで、畳み込みは、n 次元のウィンドウが n 次元のベース領域を移動し、ウィンドウの可能な位置ごとに計算が実行されるものと考えることができます。

Conv 畳み込み命令を計算にエンキューします。この計算では、拡張なしのデフォルトの畳み込みディメンション番号が使用されます。

パディングは、SAME または VALID のいずれかとして短縮形で指定されます。SAME パディングは、ストライドを考慮しない場合に出力が入力と同じ形状になるように、入力（lhs）をゼロでパディングします。VALID パディングは、単にパディングがないことを意味します。

Conv(lhs, rhs, window_strides, padding, feature_group_count, batch_group_count, precision_config, preferred_element_type)

Conv では、次のレベルの制御が可能です。

• ConvWithGeneralPadding
• ConvWithGeneralDimensions
• ConvGeneral
• ConvGeneralDilated

n を空間ディメンションの数とします。lhs 引数は、底面積を表す（n+2）次元の配列です。これは入力と呼ばれます（もちろん、右辺も入力ですが）。ニューラル ネットワークでは、これらは入力アクティベーションです。n+2 個のディメンションは次の順序で指定します。

• batch: このディメンションの各座標は、畳み込みが実行される独立した入力を表します。
• z/depth/features: ベース領域の各（y,x）位置には、このディメンションに入るベクトルが関連付けられています。
• spatial_dims: ウィンドウが移動するベース領域を定義する n 空間ディメンションを記述します。

rhs 引数は、畳み込みフィルタ/カーネル/ウィンドウを記述する（n+2）次元配列です。ディメンションは次の順序で指定します。

• output-z: 出力の z ディメンション。
• input-z: このディメンションのサイズに feature_group_count を掛けた値は、lhs の z ディメンションのサイズと等しくなるはずです。
• spatial_dims: ベースエリアを移動する n 次元ウィンドウを定義する 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 引数は、各空間次元の左辺と右辺にそれぞれ適用される拡張係数を指定します。空間次元の拡張係数が d の場合、その次元の各エントリの間に d-1 個の穴が暗黙的に配置され、配列のサイズが増加します。穴は no-op 値で埋められます。畳み込みの場合、これはゼロを意味します。

右辺の拡張は、アトラス畳み込みとも呼ばれます。詳しくは、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 を掛けた値は、lhs の batch ディメンションのサイズと等しくなるはずです。
• z: カーネルの output-z と同じサイズ（rhs）。
• spatial_dims: 畳み込みウィンドウの有効な配置ごとに 1 つの値。

上の図は、batch_group_count フィールドの仕組みを示しています。実際には、各 lhs バッチを batch_group_count グループに分割し、出力特徴についても同様の処理を行います。次に、これらの各グループに対してペアワイズ畳み込みを行い、出力特徴ディメンションに沿って出力を連結します。他のすべてのディメンション（特徴と空間）のオペレーション セマンティクスは変わりません。

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

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

output-z が 5 の場合、ウィンドウの各位置で、出力の z ディメンションに 5 つの値が出力されます。これらの値は、畳み込みカーネルのどの部分が使用されるかによって異なります。各 output-z 座標に使用される値の 3D ボックスは別個に存在します。つまり、それぞれに異なるフィルタを使用した 5 つの個別の畳み込みと考えることができます。

パディングとストライドを使用した 2D 畳み込みの擬似コードは次のとおりです。

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

precision_config は、精度構成を示すために使用されます。このレベルは、必要に応じてより正確な dtype エミュレーションを提供するために、ハードウェアがより多くのマシンコード命令を生成しようとするかどうかを決定します（bf16 matmul のみをサポートする TPU で f32 をエミュレートするなど）。値は DEFAULT、HIGH、HIGHEST のいずれかです。詳しくは、MXU セクションをご覧ください。

preferred_element_type は、累積に使用される高精度/低精度の出力型のスカラー要素です。preferred_element_type は、指定されたオペレーションの累積タイプを推奨しますが、保証はされません。これにより、一部のハードウェア バックエンドでは、別の型で累積し、優先出力型に変換できます。

StableHLO の詳細については、StableHLO - 畳み込みをご覧ください。

ConvWithGeneralPadding

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

ConvWithGeneralPadding(lhs, rhs, window_strides, padding, feature_group_count, batch_group_count, precision_config, preferred_element_type)

パディング構成が明示的な Conv と同じです。

ConvWithGeneralDimensions

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

ConvWithGeneralDimensions(lhs, rhs, window_strides, padding, dimension_numbers, feature_group_count, batch_group_count, precision_config, preferred_element_type)

ディメンション番号が明示されている Conv と同じです。

ConvGeneral

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

ConvGeneral(lhs, rhs, window_strides, padding, dimension_numbers, feature_group_count, batch_group_count, precision_config, preferred_element_type)

ディメンション番号とパディング構成が明示的な Conv と同じ

ConvGeneralDilated

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

ConvGeneralDilated(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, dimension_numbers, feature_group_count, batch_group_count, precision_config, preferred_element_type, window_reversal)

パディング構成、拡張率、ディメンション番号が明示的な Conv と同じです。

コピー

HloInstruction::CreateCopyStart もご覧ください。

Copy は、内部的に 2 つの HLO 命令 CopyStart と CopyDone に分解されます。Copy は、CopyStart、CopyDone とともに HLO のプリミティブとして機能します。これらのオペレーションは HLO ダンプに表示されることがありますが、エンドユーザーが手動で構築することを想定していません。

Cos

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

要素単位のコサイン x -> cos(x)。

Cos(operand)

Cos は、省略可能な result_accuracy 引数もサポートしています。

Cos(operand, result_accuracy)

result_accuracy の詳細については、結果の精度をご覧ください。

StableHLO の詳細については、StableHLO - cosine をご覧ください。

Cosh

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

要素ごとの双曲線余弦 x -> cosh(x)。

Cosh(operand)

Cosh は、省略可能な result_accuracy 引数もサポートしています。

Cosh(operand, result_accuracy)

result_accuracy の詳細については、結果の精度をご覧ください。

CustomCall

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

計算内でユーザーが指定した関数を呼び出します。

CustomCall のドキュメントは、デベロッパーの詳細 - XLA カスタム呼び出しで提供されています。

StableHLO の詳細については、StableHLO - custom_call をご覧ください。

Div

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

被除数 lhs と除数 rhs の要素ごとの除算を行います。

Div(lhs, rhs)

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

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

Div には、異なる次元のブロードキャストをサポートする別のバリアントがあります。

Div(lhs,rhs, broadcast_dimensions)

このオペレーションのバリアントは、ランクの異なる配列間の算術演算（行列とベクトルの加算など）に使用する必要があります。

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

StableHLO の詳細については、StableHLO - divide をご覧ください。

ドメイン

HloInstruction::CreateDomain もご覧ください。

Domain は HLO ダンプに表示されることがありますが、エンドユーザーが手動で構築することを意図したものではありません。

ドット

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

Dot(lhs, rhs, precision_config, preferred_element_type)

このオペレーションの正確なセマンティクスは、オペランドのランクによって異なります。

このオペレーションは、lhs の 2 番目のディメンション（1 つのディメンションしかない場合は 1 番目のディメンション）と rhs の 1 番目のディメンションの積の合計を計算します。これらは「契約」ディメンションです。lhs と rhs の縮約されたディメンションは同じサイズである必要があります。実際には、ベクトル間のドット積、ベクトルと行列の乗算、行列と行列の乗算を実行するために使用できます。

precision_config は、精度構成を示すために使用されます。このレベルは、必要に応じてより正確な dtype エミュレーションを提供するために、ハードウェアがより多くのマシンコード命令を生成しようとするかどうかを決定します（bf16 matmul のみをサポートする TPU で f32 をエミュレートするなど）。値は DEFAULT、HIGH、HIGHEST のいずれかです。詳しくは、MXU セクションをご覧ください。

preferred_element_type は、累積に使用される高精度/低精度の出力型のスカラー要素です。preferred_element_type は、指定されたオペレーションの累積タイプを推奨しますが、保証はされません。これにより、一部のハードウェア バックエンドでは、別の型で累積し、優先出力型に変換できます。

StableHLO の詳細については、StableHLO - dot をご覧ください。

DotGeneral

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

DotGeneral(lhs, rhs, dimension_numbers, precision_config, preferred_element_type)

Dot と同様ですが、lhs と 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、2x2 行列）:

lhs = { { {1.0, 2.0},
          {3.0, 4.0} },
        { {5.0, 6.0},
          {7.0, 8.0} } }

rhs = { { {1.0, 0.0},
          {0.0, 1.0} },
        { {1.0, 0.0},
          {0.0, 1.0} } }

DotDimensionNumbers dnums;
dnums.add_lhs_contracting_dimensions(2);
dnums.add_rhs_contracting_dimensions(1);
dnums.add_lhs_batch_dimensions(0);
dnums.add_rhs_batch_dimensions(0);

DotGeneral(lhs, rhs, dnums) -> {
    { {1.0, 2.0},
      {3.0, 4.0} },
    { {5.0, 6.0},
      {7.0, 8.0} } }

結果として得られるディメンション番号は、バッチ ディメンション、lhs 非収縮/非バッチ ディメンション、rhs 非収縮/非バッチ ディメンションの順になります。

precision_config は、精度構成を示すために使用されます。このレベルは、必要に応じてより正確な dtype エミュレーションを提供するために、ハードウェアがより多くのマシンコード命令を生成しようとするかどうかを決定します（bf16 matmul のみをサポートする TPU で f32 をエミュレートするなど）。値は DEFAULT、HIGH、HIGHEST のいずれかです。詳細については、MXU のセクションをご覧ください。

preferred_element_type は、累積に使用される高精度/低精度の出力型のスカラー要素です。preferred_element_type は、指定されたオペレーションの累積タイプを推奨しますが、保証はされません。これにより、一部のハードウェア バックエンドでは、別の型で累積し、優先出力型に変換できます。

StableHLO の詳細については、StableHLO - dot_general をご覧ください。

ScaledDot

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

ScaledDot(lhs, lhs_scale, rhs, rhs_scale, dimension_number, precision_config,preferred_element_type)

DotGeneral と同様です。

オペランド「lhs」、「lhs_scale」、「rhs」、「rhs_scale」を使用して、スケーリングされたドット演算を作成します。収縮ディメンションとバッチ ディメンションは「dimension_numbers」で指定します。

RaggedDot

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

RaggedDot 計算の内訳については、StableHLO - chlo.ragged_dot をご覧ください。

DynamicReshape

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

このオペレーションは reshape と機能的に同じですが、結果の形状は output_shape を介して動的に指定されます。

DynamicReshape(operand, dim_sizes, new_size_bounds, dims_are_dynamic)

StableHLO の詳細については、StableHLO - dynamic_reshape をご覧ください。

DynamicSlice

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

DynamicSlice は、動的 start_indices で入力配列からサブ配列を抽出します。各ディメンションのスライスのサイズは size_indices で渡されます。これは、各ディメンションの排他的スライス間隔の終了点を指定します（[start, start + size)）。start_indices の形状は 1 次元で、ディメンションのサイズは operand のディメンションの数と同じである必要があります。

DynamicSlice(operand, start_indices, slice_sizes)

有効なスライス インデックスは、スライスを実行する前に [1, N) の各インデックス i に次の変換を適用して計算されます。

start_indices[i] = clamp(start_indices[i], 0, operand.dimension_size[i] - slice_sizes[i])

これにより、抽出されたスライスは常にオペランド配列の範囲内になります。変換が適用される前にスライスが境界内にある場合、変換は効果がありません。

1 次元の例:

let a = {0.0, 1.0, 2.0, 3.0, 4.0};
let s = {2};

DynamicSlice(a, s, {2});
// Result: {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});
//Result:
// { { 7.0,  8.0},
//   {10.0, 11.0} }

StableHLO の詳細については、StableHLO - dynamic_slice をご覧ください。

DynamicUpdateSlice

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

DynamicUpdateSlice は、入力配列 operand の値である結果を生成します。この結果では、スライス update が start_indices で上書きされています。update の形状によって、更新される結果のサブ配列の形状が決まります。start_indices の形状は 1 次元で、ディメンションのサイズは operand のディメンションの数と等しくなければなりません。

DynamicUpdateSlice(operand, update, start_indices)

有効なスライス インデックスは、スライスを実行する前に [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)
// Result: {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)
// Result:
// { {0.0,  1.0,  2.0},
//   {3.0, 12.0, 13.0},
//   {6.0, 14.0, 15.0},
//   {9.0, 16.0, 17.0} }

StableHLO の詳細については、StableHLO - dynamic_update_slice をご覧ください。

Erf

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

要素ごとの誤差関数 x -> erf(x)。ここで、

\(\text{erf}(x) = \frac{2}{\sqrt{\pi} }\int_0^x e^{-t^2} \, dt\)。

Erf(operand)

Erf は、省略可能な result_accuracy 引数もサポートしています。

Erf(operand, result_accuracy)

result_accuracy の詳細については、結果の精度をご覧ください。

Exp

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

要素ごとの自然指数 x -> e^x。

Exp(operand)

Exp は、オプションの result_accuracy 引数もサポートしています。

Exp(operand, result_accuracy)

result_accuracy の詳細については、結果の精度をご覧ください。

StableHLO の詳細については、StableHLO - 指数関数をご覧ください。

Expm1

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

要素ごとの自然指数から 1 を引いた値 x -> e^x - 1。

Expm1(operand)

Expm1 は、省略可能な result_accuracy 引数もサポートしています。

Expm1(operand, result_accuracy)

result_accuracy の詳細については、結果の精度をご覧ください。

StableHLO の情報については、StableHLO - exponential_minus_one をご覧ください。

Fft

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

XLA FFT オペレーションは、実数と複素数の入力/出力に対して順方向と逆方向のフーリエ変換を実装します。最大 3 軸の多次元 FFT がサポートされています。

Fft(operand, ftt_type, fft_length)

StableHLO の詳細については、StableHLO - fft をご覧ください。

多次元 FFT

複数の fft_length が指定された場合、これは、最内側の各軸に FFT オペレーションのカスケードを適用することと同じです。実数から複素数への変換と複素数から実数への変換の場合、最も内側の軸の変換が（事実上）最初に行われます（RFFT、IRFFT の場合は最後）。そのため、最も内側の軸のサイズが変更されます。他の軸変換は複素数 -> 複素数になります。

実装の詳細

CPU FFT は Eigen の TensorFFT によってサポートされています。GPU FFT は cuFFT を使用します。

階

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

要素ごとの床関数 x -> ⌊x⌋。

Floor(operand)

StableHLO の情報については、StableHLO - floor をご覧ください。

Fusion

HloInstruction::CreateFusion もご覧ください。

Fusion オペレーションは HLO 命令を表し、HLO のプリミティブとして機能します。このオペレーションは HLO ダンプに表示されることがありますが、エンドユーザーが手動で構築することを想定していません。

収集

XLA 収集演算は、入力配列の複数のスライス（各スライスは実行時のオフセットが異なる可能性があります）を結合します。

StableHLO の詳細については、StableHLO - gather をご覧ください。

一般的なセマンティクス

XlaBuilder::Gather もご覧ください。より直感的な説明については、以下の「非公式な説明」セクションをご覧ください。

gather(operand, start_indices, dimension_numbers, slice_sizes, indices_are_sorted)

便宜上、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 に沿った出力配列の境界は、次のように計算されます。

1. 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] を選択します）。

2. i が offset_dims に存在する場合（つまり、ある k について offset_dims[k] と等しい場合）、collapsed_slice_dims を考慮したうえで、slice_sizes から対応する境界を選択します（つまり、adjusted_slice_sizes[k] を選択します。ここで、adjusted_slice_sizes は、インデックス collapsed_slice_dims の境界が削除された slice_sizes です）。

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

1. G = { Out[k] for k in batch_dims } とします。G を使用して、S[i] = start_indices[Combine(G, i)] となるベクトル S を切り出します。ここで、Combine(A, b) は、位置 index_vector_dim に b を A に挿入します。G が空の場合でも、これは適切に定義されます。G が空の場合、S = start_indices です。

2. start_index_map を使用して S を分散し、S を使用して operand に開始インデックス Sin を作成します。具体的には、次のようになります。

   1. 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. In は Oin + Sin です。ここで、+ は要素ごとの加算です。

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,O0,O1] を受け取り、出力シェイプのインデックスを次のように入力配列の要素にマッピングするインデックス変換として表すことができます。

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

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

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

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

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

XLA の gather オペレーションは、上記の非公式なセマンティクスを次のように一般化します。

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

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

3. 「収集インデックス」配列から抽出されたスライス（最後の例では（X、Y））の要素数は、入力配列のディメンション数よりも少ない場合があります。明示的なマッピングにより、入力と同じディメンション数になるようにインデックスを拡張する方法が指定されます。

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

G0 と G1 は、通常どおり、収集インデックス配列から開始インデックスを切り出すために使用されます。ただし、開始インデックスには 1 つの要素 X のみがあります。同様に、値 O0 を持つ出力オフセット インデックスは 1 つだけです。ただし、これらは入力配列のインデックスとして使用される前に、「Gather Index Mapping」（正式な説明では start_index_map）と「Offset Mapping」（正式な説明では remapped_offset_dims）に従ってそれぞれ [X,0] と [0,O0] に拡張され、[X,O0] に加算されます。つまり、出力インデックス [G0,G1,O0] は入力インデックス [GatherIndices[G0,G1,0],O0] にマッピングされ、tf.gather_nd のセマンティクスが得られます。

この場合の slice_sizes は [1,11] です。直感的には、gather インデックス配列の各インデックス X が行全体を選択し、結果はこれらのすべての行の連結になります。

GetDimensionSize

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

オペランドの指定されたディメンションのサイズを返します。オペランドは配列の形状にする必要があります。

GetDimensionSize(operand, dimension)

StableHLO の情報については、StableHLO - get_dimension_size をご覧ください。

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 もご覧ください。

GetTupleElement(tuple_data, index)

StableHLO の詳細については、StableHLO - get_tuple_element をご覧ください。

Imag

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

複素数（または実数）シェイプの要素ごとの虚数部。x -> imag(x)。オペランドが浮動小数点型の場合は、0 を返します。

Imag(operand)

StableHLO の詳細については、StableHLO - imag をご覧ください。

インフィード

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

Infeed(shape, config)

デバイスの暗黙的な Infeed ストリーミング インターフェースから単一のデータアイテムを読み取り、データを指定された形状とそのレイアウトとして解釈して、データの XlaOp を返します。計算では複数の Infeed オペレーションが許可されますが、Infeed オペレーション間には全順序が必要です。たとえば、次のコードの 2 つの Infeed は、while ループ間に依存関係があるため、全順序になります。

result1 = while (condition, init = init_value) {
  Infeed(shape)
  }

result2 = while (condition, init = result1) {
  Infeed(shape)
  }

ネストされたタプル形状はサポートされていません。空のタプル形状の場合、Infeed オペレーションは事実上 no-op であり、デバイスの Infeed からデータを読み取らずに続行します。

StableHLO の詳細については、StableHLO - インフィードをご覧ください。

Iota

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

Iota(shape, iota_dimension)

潜在的に大きなホスト転送ではなく、デバイス上に定数リテラルを構築します。指定された形状を持ち、指定されたディメンションに沿って 0 から 1 ずつ増加する値を保持する配列を作成します。浮動小数点型の場合、生成される配列は ConvertElementType(Iota(...)) と同等です。ここで、Iota は整数型で、変換は浮動小数点型です。

たとえば、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 ]]

StableHLO の情報については、StableHLO - iota をご覧ください。

IsFinite

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

operand の各要素が有限であるかどうか（正または負の無限大ではなく、NaN でもないかどうか）をテストします。入力と同じ形状の PRED 値の配列を返します。各要素は、対応する入力要素が有限である場合にのみ true になります。

IsFinite(operand)

StableHLO の情報については、StableHLO - is_finite をご覧ください。

ログ

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

要素ごとの自然対数 x -> ln(x)。

Log(operand)

Log は、オプションの result_accuracy 引数もサポートしています。

Log(operand, result_accuracy)

result_accuracy の詳細については、結果の精度をご覧ください。

StableHLO の情報については、StableHLO - ログをご覧ください。

Log1p

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

要素ごとのシフトされた自然対数 x -> ln(1+x)。

Log1p(operand)

Log1p は、オプションの result_accuracy 引数もサポートしています。

Log1p(operand, result_accuracy)

result_accuracy の詳細については、結果の精度をご覧ください。

StableHLO の詳細については、StableHLO - log_plus_one をご覧ください。

ロジスティック

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

要素ごとのロジスティック関数計算 x -> logistic(x)。

Logistic(operand)

Logistic は、オプションの result_accuracy 引数もサポートしています。

Logistic(operand, result_accuracy)

result_accuracy の詳細については、結果の精度をご覧ください。

StableHLO の詳細については、StableHLO - ロジスティックをご覧ください。

地図

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

Map(operands..., computation, dimensions)

指定された operands 配列にスカラー関数を適用し、同じディメンションの配列を生成します。この配列の各要素は、入力配列の対応する要素に適用されたマッピング関数の結果です。

マッピングされた関数は任意の計算ですが、スカラー型 T の N 個の入力と、型 S の単一の出力があるという制限があります。出力のディメンションはオペランドと同じですが、要素型 T が S に置き換えられます。

たとえば、Map(op1, op2, op3, computation, par1) は、入力配列の各（多次元）インデックスで elem_out <- computation(elem1, elem2, elem3, par1) をマッピングして、出力配列を生成します。

StableHLO の詳細については、StableHLO - map をご覧ください。

最大

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

テンソル lhs と rhs に対して要素ごとの最大演算を実行します。

Max(lhs, rhs)

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

Max には、異なる次元のブロードキャストをサポートする別のバリアントがあります。

Max(lhs,rhs, broadcast_dimensions)

このオペレーションのバリアントは、ランクの異なる配列間の算術演算（行列とベクトルの加算など）に使用する必要があります。

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

StableHLO の情報については、StableHLO - 最大をご覧ください。

最小

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

lhs と rhs で要素ごとの最小値演算を実行します。

Min(lhs, rhs)

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

Min には、異なる次元のブロードキャストをサポートする代替バリアントがあります。

Min(lhs,rhs, broadcast_dimensions)

このオペレーションのバリアントは、ランクの異なる配列間の算術演算（行列とベクトルの加算など）に使用する必要があります。

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

StableHLO の情報については、StableHLO - 最小をご覧ください。

Mul

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

lhs と rhs の要素ごとの積を計算します。

Mul(lhs, rhs)

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

Mul には、異なる次元のブロードキャストをサポートする別のバリアントがあります。

Mul(lhs,rhs, broadcast_dimensions)

このオペレーションのバリアントは、ランクの異なる配列間の算術演算（行列とベクトルの加算など）に使用する必要があります。

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

StableHLO の情報については、StableHLO - multiply をご覧ください。

Neg

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

要素ごとの否定 x -> -x。

Neg(operand)

StableHLO の情報については、StableHLO - negate をご覧ください。

×

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

要素ごとの論理否定 x -> !(x)。

Not(operand)

StableHLO の情報については、StableHLO - not をご覧ください。

OptimizationBarrier

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

最適化パスがバリアを越えて計算を移動することをブロックします。

OptimizationBarrier(operand)

バリアの出力に依存する演算子の前に、すべての入力が評価されるようにします。

StableHLO の詳細については、StableHLO - optimization_barrier をご覧ください。

または

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

lhs と rhs の要素ごとの OR を実行します。

Or(lhs, rhs)

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

Or には、異なる次元のブロードキャストをサポートする別のバリエーションがあります。

Or(lhs,rhs, broadcast_dimensions)

このオペレーションのバリアントは、ランクの異なる配列間の算術演算（行列とベクトルの加算など）に使用する必要があります。

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

StableHLO の詳細については、StableHLO - またはをご覧ください。

アウトフィード

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

入力をアウトフィードに書き込みます。

Outfeed(operand, shape_with_layout, outfeed_config)

shape_with_layout は、アウトフィードするレイアウトされた形状を伝えます。

StableHLO の情報については、StableHLO - アウトフィードをご覧ください。

パッド

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

Pad(operand, padding_value, padding_config)

指定された 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 の場合、このオペレーションは no-op です。次の図は、2 次元配列のさまざまな edge_padding 値と interior_padding 値の例を示しています。

StableHLO の詳細については、StableHLO - pad をご覧ください。

パラメータ

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

Parameter は、計算への引数入力を表します。

PartitionID

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

現在のプロセスの partition_id を生成します。

PartitionID(shape)

PartitionID は HLO ダンプに表示されることがありますが、エンドユーザーが手動で構築することを意図したものではありません。

StableHLO の情報については、StableHLO - partition_id をご覧ください。

PopulationCount

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

operand の各要素に設定されているビット数を計算します。

PopulationCount(operand)

StableHLO の詳細については、StableHLO - popcnt をご覧ください。

Pow

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

lhs の要素と rhs の要素の累乗を計算します。

Pow(lhs, rhs)

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

Pow には、異なる次元のブロードキャストをサポートする別のバリアントがあります。

Pow(lhs,rhs, broadcast_dimensions)

このオペレーションのバリアントは、ランクの異なる配列間の算術演算（行列とベクトルの加算など）に使用する必要があります。

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

StableHLO の詳細については、StableHLO - power をご覧ください。

実際

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

複素数（または実数）シェイプの要素ごとの実部。x -> real(x)。オペランドが浮動小数点型の場合、Real は同じ値を返します。

Real(operand)

StableHLO の情報については、StableHLO - real をご覧ください。

Recv

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

Recv、RecvWithTokens、RecvToHost は、HLO で通信プリミティブとして機能するオペレーションです。これらのオペレーションは通常、低レベルの入出力またはデバイス間の転送の一部として HLO ダンプに表示されますが、エンドユーザーが手動で構築することを想定していません。

Recv(shape, handle)

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

StableHLO の詳細については、StableHLO - recv をご覧ください。

RecvDone

HloInstruction::CreateRecv と HloInstruction::CreateRecvDone もご覧ください。

Send と同様に、Recv オペレーションのクライアント API は同期通信を表します。ただし、この命令は内部的に 2 つの HLO 命令（Recv と RecvDone）に分解され、非同期データ転送が可能になります。

Recv(const Shape& shape, int64 channel_id)

同じ channel_id の Send 命令からデータを受信するために必要なリソースを割り当てます。割り当てられたリソースのコンテキストを返します。このコンテキストは、後続の RecvDone 命令でデータ転送の完了を待つために使用されます。コンテキストは {受信バッファ（形状）、リクエスト ID（U32）} のタプルであり、RecvDone 命令でのみ使用できます。

Recv 命令で作成されたコンテキストを指定して、データ転送が完了するまで待機し、受信したデータを返します。

削減

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

1 つ以上の配列にリダクション関数を並行して適用します。

Reduce(operands..., init_values..., computation, dimensions_to_reduce)

ここで

• N は 1 以上である必要があります。
• 計算は「ほぼ」結合的である必要があります（下記を参照）。
• 入力配列はすべて同じディメンションである必要があります。
• すべての初期値は、computation の下で ID を形成する必要があります。
• N = 1 の場合、Collate(T) は T です。
• N > 1、Collate(T_0, ..., T_{N-1}) が T 型の N 要素のタプルの場合。

このオペレーションは、各入力配列の 1 つ以上のディメンションをスカラーに縮小します。返される各配列のディメンション数は number_of_dimensions(operand) - len(dimensions) です。このオペレーションの出力は Collate(Q_0, ..., Q_N) です。ここで、Q_i は T_i 型の配列です。そのディメンションについては後述します。

異なるバックエンドで削減計算を再関連付けることができます。加算などの一部の削減関数は浮動小数点数に対して結合的ではないため、数値の差異が生じる可能性があります。ただし、データの範囲が制限されている場合、浮動小数点加算はほとんどの実用的な用途で結合法則を満たすほど十分に近似しています。

StableHLO の詳細については、StableHLO - reduce をご覧ください。

例

値 [10, 11, 12, 13] を含む 1 つの 1 次元配列の 1 つのディメンションで削減する場合、削減関数 f（これは computation）を使用すると、次のように計算できます。

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 number of dimensions 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 で、次元 0 のサイズは 4、次元 1 のサイズは 2、次元 2 のサイズは 3 です。簡単にするため、値 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 を追加すると、1D 配列 [20, 28, 36] が生成されます。

3 次元配列のすべての次元を縮小すると、スカラー 84 が生成されます。

Variadic Reduce

N > 1 の場合、関数適用はすべての入力に同時に適用されるため、やや複雑になります。オペランドは次の順序で計算に渡されます。

• 最初のオペランドの実行値の削減
• ...
• N 番目のオペランドの値を減らして実行
• 最初のオペランドの入力値
• ...
• N 番目のオペランドの入力値

たとえば、次の削減関数は、1 次元配列の最大値と 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] と初期値 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, exponent_bits, mantissa_bits)

結果は T 型の配列です。入力値は、指定された仮数部のビット数で表現可能な最も近い値に丸められ（「偶数への丸め」セマンティクスを使用）、指数部のビット数で指定された範囲を超える値は正または負の無限大にクランプされます。NaN 値は保持されますが、正規の NaN 値に変換される場合があります。

低精度形式には、少なくとも 1 つの指数ビットが必要です（ゼロ値と無限大を区別するため。どちらも仮数部はゼロです）。また、仮数部のビット数は 0 以上である必要があります。指数または仮数のビット数が型 T の対応する値を超えることがあります。その場合、変換の対応する部分は単に no-op になります。

StableHLO の情報については、StableHLO - reduce_precision をご覧ください。

ReduceScatter

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

ReduceScatter は、AllReduce を実行してから、結果を scatter_dimension に沿って shard_count ブロックに分割して分散する集団演算です。レプリカ グループのレプリカ i は ith シャードを受け取ります。

ReduceScatter(operand, computation, scatter_dimension, shard_count, replica_groups, channel_id, layout, use_global_device_ids)

• operand が配列のタプルの場合、タプルの各要素に対して reduce-scatter が実行されます。
• replica_groups は、縮小が実行されるレプリカ グループのリストです（現在のレプリカのレプリカ ID は ReplicaId を使用して取得できます）。各グループ内のレプリカの順序によって、all-reduce の結果が分散される順序が決まります。replica_groups は空（すべてのレプリカが単一のグループに属している場合）であるか、レプリカの数と同じ数の要素を含んでいる必要があります。レプリカ グループが複数ある場合は、すべて同じサイズにする必要があります。たとえば、replica_groups = {0, 2}, {1, 3} はレプリカ 0 と 2、1 と 3 の間で削減を実行し、結果を分散します。
• shard_count は各レプリカ グループのサイズです。これは、replica_groups が空の場合に必要です。replica_groups が空でない場合、shard_count は各レプリカ グループのサイズと等しくなければなりません。
• channel_id はモジュール間の通信に使用されます。同じ channel_id を持つ reduce-scatter オペレーションのみが相互に通信できます。
• layout レイアウトについて詳しくは、xla::shapes をご覧ください。
• use_global_device_ids は、ユーザーが指定したフラグです。false（デフォルト）の場合、replica_groups の数値は ReplicaId です。true の場合、replica_groups は（ReplicaID*partition_count + partition_id）のグローバル ID を表します。例:
  • 2 つのレプリカと 4 つのパーティションがある場合、
  • replica_groups={ {0,1,4,5},{2,3,6,7} } and use_global_device_ids=true
  • group[0] = (0,0), (0,1), (1,0), (1,1)
  • group[1] = (0,2), (0,3), (1,2), (1,3)
  • 各ペアは（replica_id、partition_id）です。

出力形状は、scatter_dimension が shard_count 倍小さくなった入力形状です。たとえば、2 つのレプリカがあり、オペランドの値が 2 つのレプリカでそれぞれ [1.0, 2.25] と [3.0, 5.25] の場合、scatter_dim が 0 のこの op からの出力値は、最初のレプリカでは [4.0]、2 番目のレプリカでは [7.5] になります。

StableHLO の詳細については、StableHLO - reduce_scatter をご覧ください。

ReduceScatter - 例 1 - StableHLO

上記の例では、ReduceScatter に 2 つのレプリカが参加しています。各レプリカでは、オペランドの形状は f32[2,4] です。レプリカ間で all-reduce（合計）が実行され、各レプリカで f32[2,4] 形状の縮小値が生成されます。この縮小された値は、ディメンション 1 に沿って 2 つの部分に分割されるため、各部分の形状は f32[2,2] になります。プロセス グループ内の各レプリカは、グループ内の位置に対応する部分を受け取ります。その結果、各レプリカの出力の形状は f32[2,2] になります。

ReduceWindow

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

N 個の多次元配列のシーケンスの各ウィンドウ内のすべての要素に削減関数を適用し、出力として 1 つまたは N 個の多次元配列のタプルを生成します。各出力配列には、ウィンドウの有効な位置の数と同じ数の要素があります。プーリング レイヤは ReduceWindow として表すことができます。Reduce と同様に、適用された computation には常に左側の init_values が渡されます。

ReduceWindow(operands..., init_values..., computation, window_dimensions, window_strides, padding)

ここで

• N は 1 以上である必要があります。
• 入力配列はすべて同じディメンションである必要があります。
• N = 1 の場合、Collate(T) は T です。
• N > 1、Collate(T_0, ..., T_{N-1}) が (T0,...T{N-1}) 型の N 要素のタプルの場合。

StableHLO の情報については、StableHLO - reduce_window をご覧ください。

ReduceWindow - 例 1

入力は [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);