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)

參數:

• 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.

回傳:

The output of the softmax.

回傳型別:

array