eazygrad.softmax¶

eazygrad.softmax(input: _Tensor, dim: int) → _Tensor[source]¶

Compute the softmax of a tensor along a given axis.

Parameters:

input (_Tensor) – Input tensor.
dim (int) – Axis along which the softmax is computed.

Returns:

_Tensor – Tensor of normalized exponentials.

See also