TPU レプリカ間でデータを交換する Op。
各レプリカでは、入力が `split_dimension` に沿って `split_count` ブロックに分割され、group_assignment を指定して他のレプリカに送信されます。他のレプリカから `split_count` - 1 ブロックを受信した後、出力として `concat_dimension` に沿ってブロックを連結します。
たとえば、2 つの TPU レプリカがあるとします。レプリカ 0 は入力を受け取ります: `[[A, B]]` レプリカ 1 は入力を受け取ります: `[[C, D]]`
group_assignment=`[[0, 1]]` concat_dimension=0 split_dimension=1 split_count=2
レプリカ 0 の出力: `[[A], [C]]` レプリカ 1 の出力: `[[B], [D]]`
パブリックメソッド
public Output <T> asOutput ()
テンソルのシンボリック ハンドルを返します。
TensorFlow オペレーションへの入力は、別の TensorFlow オペレーションの出力です。このメソッドは、入力の計算を表すシンボリック ハンドルを取得するために使用されます。
public static AllToAll <T> create (スコープscope、オペランド<T>入力、オペランド<Integer>groupAssignment、Long concatDimension、Long splitDimension、Long splitCount)
新しい AllToAll 操作をラップするクラスを作成するファクトリ メソッド。
パラメーター
範囲 | 現在のスコープ |
---|---|
入力 | 合計へのローカル入力。 |
グループ割り当て | 形状 [num_groups, num_replicas_per_group] の int32 テンソル。 「group_assignment[i]」は、i 番目のサブグループ内のレプリカ ID を表します。 |
連結次元 | 連結する次元番号。 |
分割ディメンション | 分割する次元の番号。 |
分割数 | 分割の数。この数はサブグループ サイズ(group_assignment.get_shape()[1]) と等しくなければなりません。 |
戻り値
- AllToAll の新しいインスタンス