linalg:softmax

(linalg:softmax array &key axis)

array の softmax、すなわち合計が 1 になるよう正規化した exp(x - max) を返します。:axis なしでは配列全体が 1 つの分布になり (scipy の softmax のデフォルト)、整数の :axis (負の値は末尾から数えます) を渡すとその軸に沿った各スライスが個別に正規化されます。後者が torch の softmax(x, dim) に相当するアテンション重みの形です。最大値を先に引くため大きなロジットでもオーバーフローせず、-infinity の要素 (マスクされた位置、linalg:where を参照) はちょうど 0.0 になります。

linalg:relu と同様に softmax は numpy 本体にはありませんが、活性化層が必要とする配列レベルのプリミティブなのでここに置いています。対数版は linalg:log-softmax です。