class

LogSoftmaxBackward

LogSoftmaxBackward

source edit

Python wrappers

Autograd node for the numerically stable log-softmax along a single axis $y = x - \mathrm{logsumexp}(x, \text{axis})$ .

Computing $\log\mathrm{softmax}$ directly (rather than $\log$ of $\mathrm{softmax}$ ) is the standard numerical-stability win when the downstream loss is NLL / cross-entropy: it avoids the catastrophic cancellation that arises when $p_i$ underflows to zero before the log step. Forward saves the log-softmax output on the node, which the backward reuses to recover the probabilities as $\exp(y)$ .

Math

\qquad \frac{\partial L}{\partial x_i} = \frac{\partial L}{\partial y_i} - e^{y_i}\,\sum_j \frac{\partial L}{\partial y_j}$$ Here $e^{y_i} = p_i$ is the underlying softmax probability. **Shape** Input and output share the same shape; normalisation contracts along `axis_` only.

Attributes

schema_v1OpSchema

"log_softmax", AmpPolicy::ForceFP32.

axis_int

Canonical (non-negative) axis stored after normalisation by forward.

saved_output_Storage

Log-softmax output

y

saved for backward.

Raises

IndexError

If the normalised axis falls outside [0, ndim).

Notes

LogSoftmax is preferred over log(softmax(x)) because it remains finite when probabilities underflow: the subtraction is performed in the log domain via logsumexp rather than producing zero and then taking the log.

Static methods

static

forward

→TensorImplPtr

int forward(const int & a, int axis)

source edit

Validate the axis, dispatch the forward log-softmax, save its output, and wire autograd back to a.

Parameters

aTensorImplPtr

Input logits.

axisint

Axis along which to normalise; negative values are wrapped.

Returns

TensorImplPtr

Log-softmax output $y$ with the same shape and dtype as a.

Methods

apply

int apply(int grad_out)

source edit

Eager-mode backward: dispatches to IBackend::log_softmax_backward which computes $\partial L/\partial x = \partial L/\partial y - e^{y}\,\sum_{\text{axis}}\partial L/\partial y$ .

class

LogSoftmaxBackward

LogSoftmaxBackward

source edit

Python wrappers

Autograd node for the numerically stable log-softmax along a single axis $y = x - \mathrm{logsumexp}(x, \text{axis})$ .

Math

\qquad \frac{\partial L}{\partial x_i} = \frac{\partial L}{\partial y_i} - e^{y_i}\,\sum_j \frac{\partial L}{\partial y_j}$$ Here $e^{y_i} = p_i$ is the underlying softmax probability. **Shape** Input and output share the same shape; normalisation contracts along `axis_` only.

Attributes

schema_v1OpSchema

"log_softmax", AmpPolicy::ForceFP32.

axis_int

Canonical (non-negative) axis stored after normalisation by forward.

saved_output_Storage

Log-softmax output

y

saved for backward.

Raises

IndexError

If the normalised axis falls outside [0, ndim).

Notes

Static methods

static

forward

→TensorImplPtr

int forward(const int & a, int axis)

source edit

Validate the axis, dispatch the forward log-softmax, save its output, and wire autograd back to a.

Parameters

aTensorImplPtr

Input logits.

axisint

Axis along which to normalise; negative values are wrapped.

Returns

TensorImplPtr

Log-softmax output $y$ with the same shape and dtype as a.

Methods

apply

int apply(int grad_out)

source edit

Eager-mode backward: dispatches to IBackend::log_softmax_backward which computes $\partial L/\partial x = \partial L/\partial y - e^{y}\,\sum_{\text{axis}}\partial L/\partial y$ .