Python wrappers
Public Python APIs implemented by this engine symbols.Autograd node for the numerically stable softmax along a single axis with .
Subtracting the per-row maximum before exponentiating is the standard
log-sum-exp trick — it preserves the analytic result while keeping every
exponentiated term in , which is essential under ForceFP32
AMP semantics for fp16/bf16 logits. The softmax probabilities are
saved on the node so the backward needs no re-evaluation.
Math
\qquad \frac{\partial L}{\partial x_i} = p_i\,\Bigl(\frac{\partial L}{\partial y_i} - \sum_j p_j\,\frac{\partial L}{\partial y_j}\Bigr)$$ The Jacobian-vector form sidesteps materialising the dense $N \times N$ softmax Jacobian and runs in $O(N)$ per row. **Shape** Input and output share the same shape; the normalisation contracts only along `axis_`. Negative axis values are normalised to non-negative indices by `forward` before being stored. **References** Bishop, "Pattern Recognition and Machine Learning", §4.3.4.Attributes
schema_v1OpSchema"softmax", AmpPolicy::ForceFP32 (probability underflow on float16 is unsafe for large logit magnitudes).axis_intCanonical (non-negative) axis along which the softmax is computed.
saved_output_StorageSoftmax probabilities produced by
forward, used by apply.Raises
IndexErrorIf the normalised axis falls outside
[0, ndim).Static methods
1Validate the axis, dispatch the forward softmax through the backend, save the resulting probabilities on the node, and wire the autograd edges back to a.
Parameters
aTensorImplPtrInput logits.
axisintSoftmax axis; negative values are normalised to
axis + ndim.Returns
TensorImplPtrOutput probabilities with the same shape and dtype as a.
Raises
IndexErrorIf
axis is out of range after normalisation.