mlx.core.softmax

Contents

mlx.core.softmax#

softmax(a: array, /, axis: None | int | Sequence[int] = None, *, precise: bool = False, stream: StreamOrDevice = 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.

  • precise (bool, optional) – Accumulate in float32 for inputs of lower precision. Otherwise the accumulation type matches the input, which can lose precision over long reduction axes. Default: False.

Returns:

The output of the softmax.

Return type:

array

Contents