mlx.core.softmax

Contents

mlx.core.softmax#

softmax(a: array, /, axis: None | int | Sequence[int] = None, *, stream: None | Stream | Device = None) → array#

Perform the softmax along the given axis.

This operation is a numerically stable version of:

exp(a) / sum(exp(a), axis, keepdims=True)
Parameters:

  • a (array) – Input array.

  • axis (int or list(int), optional) – Optional axis or axes to compute the softmax over. If unspecified this performs the softmax over the full array.

Returns:

The output of the softmax.

Return type:

array

Contents