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
