cleanlab--cleanlab
Elías Snorrason
4ce9f771c8
* unify softmax implementations * Refine mypy type hint in _set_fine_search_range function - Update type hint for `min_entropy_ind` from built-in `int` to `np.intp`. - This refinement addresses a type compatibility warning. - `np.intp` is the integer type used by numpy for indexing and can differ in size from the built-in Python `int` depending on the platform (32-bit vs 64-bit). - Mypy highlighted this type hint discrepancy. * add unit tests for softmax --------- Co-authored-by: Hui Wen <45724323+huiwengoh@users.noreply.github.com>
38 行
1.0 KiB
Python
38 行
1.0 KiB
Python
from typing import Optional
|
|
import numpy as np
|
|
|
|
|
|
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 / temperature
|
|
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)
|