wehub-resource-sync
72 lines
2.5 KiB
Python
72 lines
2.5 KiB
Python
import pytest
|
|
import numpy as np
|
|
|
|
from cleanlab.internal.numerics import softmax
|
|
|
|
|
|
class TestSoftmax:
|
|
def test_basic_softmax(self):
|
|
input_arr = np.array([1.0, 2.0, 3.0])
|
|
output = softmax(input_arr)
|
|
expected_output = np.array([0.09003057, 0.24472847, 0.66524096])
|
|
assert np.isclose(np.sum(output), 1.0)
|
|
assert np.allclose(output, expected_output)
|
|
|
|
def test_temperature_effect(self):
|
|
input_arr = np.array([1.0, 2.0, 3.0])
|
|
output_high_temp = softmax(input_arr, temperature=5.0)
|
|
output_low_temp = softmax(input_arr, temperature=0.1)
|
|
|
|
expected_high_temp = np.array([0.2693075, 0.32893292, 0.40175958])
|
|
expected_low_temp = np.array([2.06106005e-09, 4.53978686e-05, 9.99954600e-01])
|
|
|
|
assert np.allclose(output_high_temp, expected_high_temp)
|
|
assert np.allclose(output_low_temp, expected_low_temp)
|
|
|
|
def test_axis(self):
|
|
input_arr = np.array(
|
|
[
|
|
[1, 2, 3], # unit step
|
|
[4, 5, 6], # unit step
|
|
[7, 8, 10], # non-unit step
|
|
]
|
|
)
|
|
output = softmax(input_arr, axis=1)
|
|
|
|
expected_output = np.array(
|
|
[
|
|
[0.09003057, 0.24472847, 0.66524096], # unit step
|
|
[0.09003057, 0.24472847, 0.66524096], # unit step
|
|
[0.04201007, 0.1141952, 0.84379473], # non-unit step
|
|
]
|
|
)
|
|
|
|
assert np.allclose(output, expected_output)
|
|
|
|
@pytest.mark.parametrize(
|
|
"input_arr, expected_output",
|
|
[
|
|
(np.array([1.0, 2.0, 3.0]) + 1000, np.array([0.09003057, 0.24472847, 0.66524096])),
|
|
(np.array([1e3, 2e3, 3e3]), np.array([0, 0, 1])),
|
|
],
|
|
)
|
|
def test_shift(self, input_arr, expected_output):
|
|
# Without shift, softmax overflows and gets a RuntimeWarning, but just returns nan
|
|
with pytest.warns(RuntimeWarning):
|
|
output_no_shift = softmax(input_arr, shift=False)
|
|
assert np.isnan(output_no_shift).all()
|
|
|
|
output_shift = softmax(input_arr, shift=True)
|
|
assert np.allclose(output_shift, expected_output)
|
|
|
|
@pytest.mark.parametrize(
|
|
"input_arr, expected_output",
|
|
[
|
|
(np.array([0, -np.inf, -np.inf]), np.array([1.0, 0.0, 0.0])),
|
|
(np.array([-np.inf, 0, 1]), np.array([0.0, 0.26894142, 0.73105858])),
|
|
],
|
|
)
|
|
def test_special_values(self, input_arr, expected_output):
|
|
output = softmax(input_arr)
|
|
assert np.allclose(output, expected_output)
|