torch.sparse.softmax — PyTorch 2.7 documentation (original) (raw)

torch.sparse.softmax(input, dim, *, dtype=None) → Tensor

Applies a softmax function.

Softmax is defined as:

Softmax(xi)=exp(xi)∑jexp(xj)\text{Softmax}(x_{i}) = \frac{exp(x_i)}{\sum_j exp(x_j)}

where i,ji, j run over sparse tensor indices and unspecified entries are ignores. This is equivalent to defining unspecified entries as negative infinity so that exp(xk)=0exp(x_k) = 0 when the entry with index kk has not specified.

It is applied to all slices along dim, and will re-scale them so that the elements lie in the range [0, 1] and sum to 1.

Parameters