# mlx.core.softmax — MLX 0.0.7 documentation

Skip to main content
`Ctrl`+`K`

# mlx.core.softmax

mlx.core.softmax(a: array, /, axis: Union[ None , int , List[ int ]] = None, *, stream: Union[ 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