softmax

→Tensor

softmax(input: Tensor, dim: _int = ..., dtype: DTypeLike | None = ...)

source edit

Implementing kernel

Apply the softmax function along a dimension.

Parameters

inputTensor

Input tensor of any shape (the "logits").

dimint or None

Dimension along which softmax is computed.

_stacklevelint

Internal stack level for the implicit-dimension warning.

dtypeDTypeLike or None

Target floating-point dtype to cast input to before softmax.

Returns

Tensor

Same-shape tensor whose entries along dim sum to 1.

Notes

\text{softmax}(x)_i = \frac{e^{x_i}}{\sum_j e^{x_j}}

The engine uses the log-sum-exp shift $x_i \mapsto x_i - \max_j x_j$ . For loss computation, prefer log_softmax followed by nll_loss.

Examples

>>> import lucid
>>> from lucid.nn.functional import softmax
>>> logits = lucid.tensor([[1.0, 2.0, 3.0]])
>>> p = softmax(logits, dim=1)
>>> p
tensor([[0.09003, 0.2447, 0.6652]])
>>> p.sum(dim=1)
tensor([1.])

>>> import lucid >>> from lucid.nn.functional import softmax >>> logits = lucid.tensor([[1.0, 2.0, 3.0]]) >>> p = softmax(logits, dim=1) >>> p tensor([[0.09003, 0.2447, 0.6652]]) >>> p.sum(dim=1) tensor([1.])