linalg:matmul

(linalg:matmul a b)

a と b の行列積です (numpy の np.matmul、@ 演算子)。rank が 2 以下の場合は linalg:dot と同様に振る舞いますが (行列 . ベクタも可能です)、いずれかのオペランドがスカラーの場合はエラーを通知します。これにより、要素ごとの linalg:mul が暗黙に適用されてしまう場面で行列積を書いてしまう間違いを検出できます。内側の次元は一致しなければならず、不一致の場合はエラーを通知します。

いずれかの rank が 3 以上の場合はスタックされた行列積になります (torch の bmm / matmul)。最後の 2 軸が行列で、先頭側のすべての軸は numpy の規則でブロードキャストされます。したがって (batch heads n d) の query と (batch heads d n) の key の積は (batch heads n n) のアテンションスコアになります。rank 1 のオペランドは積のために昇格され (左側なら行、右側なら列)、結果からはその軸が再び落とされます。numpy とまったく同じ挙動です。