torch:softmax

(torch:softmax a &key axis)

Differentiable max-subtracted softmax (linalg:softmax): with no :axis the whole tensor is one distribution, with an integer :axis one distribution per slice -- torch's softmax(x, dim), the attention-weight form. The backward pass is s * (g - sum(g * s)) over each distribution. Masked positions filled with -infinity by torch:masked-fill come out as exactly 0.0.