Files
2026-07-13 12:49:22 +08:00

40 lines
1.1 KiB
Python

from typing import Optional
import numpy as np
from cleanlab.internal.constants import EPSILON
def softmax(
x: np.ndarray, temperature: float = 1.0, axis: Optional[int] = None, shift: bool = False
) -> np.ndarray:
"""Softmax function.
Parameters
----------
x : np.ndarray
Input array.
temperature : float
Temperature of the softmax function.
axis : Optional[int]
Axis to apply the softmax function. If None, the softmax function is
applied to all elements of the input array.
shift : bool
Whether to shift the input array before applying the softmax function.
This is useful to avoid numerical issues when the input array contains
large values, that could result in overflows when applying the exponential
function.
Returns
-------
np.ndarray
Softmax function applied to the input array.
"""
x = x / max(temperature, EPSILON)
if shift:
x = x - np.max(x, axis=axis, keepdims=True)
exp_x = np.exp(x)
return exp_x / np.sum(exp_x, axis=axis, keepdims=True)