torch:gather

(torch:gather a idx)

行列の微分可能な行ごとの選択です。各行 i について要素 a[i, idx[i]] をベクトルとして返します (linalg:gather -- 交差エントロピー損失の「正解ロジットを選ぶ」イディオム)。idx はインデックスベクトル、リスト、テンソルのいずれでも構いません。backward は選ばれたセルに勾配を散布して戻します。