oneflow.nn.functional.softmax¶
-
oneflow.nn.functional.
softmax
(input: oneflow.Tensor, dim: Optional[int] = None, dtype=None) → oneflow.Tensor¶ Applies a softmax function. Softmax is defined as: \(\text{Softmax}(x_{i}) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}\) 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.
See
Softmax
for more details.- Parameters
input (Tensor) – input
dim (int) – A dimension along which softmax will be computed.
dtype (
oneflow.dtype
, optional) – the desired data type of returned tensor. If specified, the input tensor is casted todtype
before the operation is performed. This is useful for preventing data type overflows. Default: None.
Note
This function doesn’t work directly with NLLLoss, which expects the Log to be computed between the Softmax and itself. Use log_softmax instead (it’s faster and has better numerical properties).