pytorch / 2 / generated / torch.sparse.softmax.html

torch.sparse.softmax

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

Applies a softmax function.

Softmax is defined as:

Softmax ( x i ) = e x p ( x i ) j e x p ( x j ) \text{Softmax}(x_{i}) = \frac{exp(x_i)}{\sum_j exp(x_j)}

where i , j i, j run over sparse tensor indices and unspecified entries are ignores. This is equivalent to defining unspecified entries as negative infinity so that e x p ( x k ) = 0 exp(x_k) = 0 when the entry with index k k 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
  • input (Tensor) – input
  • dim (int) – A dimension along which softmax will be computed.
  • dtype (torch.dtype, optional) – the desired data type of returned tensor. If specified, the input tensor is casted to dtype before the operation is performed. This is useful for preventing data type overflows. Default: None

© 2024, PyTorch Contributors
PyTorch has a BSD-style license, as found in the LICENSE file.
https://pytorch.org/docs/2.1/generated/torch.sparse.softmax.html