✨ special: improved [log_]softmax annotations

Author: jorenham

scipy-stubs/special/_logsumexp.pyi

@@ -11,6 +11,7 @@ __all__ = ["log_softmax", "logsumexp", "softmax"]
 _InexactT = TypeVar("_InexactT", bound=npc.inexact)
 _FloatingT = TypeVar("_FloatingT", bound=npc.floating)
 _CFloatingT = TypeVar("_CFloatingT", bound=npc.complexfloating)
+_InexactOrArrayT = TypeVar("_InexactOrArrayT", bound=npc.inexact | onp.ArrayND[npc.inexact])
 
 ###
 
@@ -98,7 +99,7 @@ def logsumexp(
     keepdims: bool = False,
     return_sign: Falsy = False,
 ) -> onp.ArrayND[np.float64 | Any] | Any: ...
-@overload  # ccomplex fallback, return_sign=False
+@overload  # complex fallback, return_sign=False
 def logsumexp(
     a: onp.ToComplex | onp.ToComplexND,
     axis: AnyShape | None = None,
@@ -223,7 +224,7 @@ def logsumexp(
     *,
     return_sign: Truthy,
 ) -> tuple[onp.ArrayND[np.float64 | Any] | Any, onp.ArrayND[np.float64 | Any] | Any]: ...
-@overload  # ccomplex fallback, return_sign=True
+@overload  # complex fallback, return_sign=True
 def logsumexp(
     a: onp.ToComplex | onp.ToComplexND,
     axis: AnyShape | None = None,
@@ -233,22 +234,46 @@ def logsumexp(
     return_sign: Truthy,
 ) -> tuple[onp.ArrayND[np.float64 | Any] | Any, onp.ArrayND[np.complex128 | Any] | Any]: ...
 
-#
-@overload
-def softmax(x: onp.ToFloat, axis: AnyShape | None = None) -> np.float64: ...
-@overload
-def softmax(x: onp.ToFloatND, axis: AnyShape | None = None) -> onp.ArrayND[np.float64]: ...
-@overload
-def softmax(x: onp.ToComplex, axis: AnyShape | None = None) -> np.float64 | np.complex128: ...
-@overload
-def softmax(x: onp.ToComplexND, axis: AnyShape | None = None) -> onp.ArrayND[np.float64 | np.complex128]: ...
+# NOTE: keep in sync with `log_softmax`
+@overload  # T
+def softmax(x: _InexactOrArrayT, axis: AnyShape | None = None) -> _InexactOrArrayT: ...  # type: ignore[overload-overlap]
+@overload  # 0d +float64
+def softmax(x: onp.ToInt | onp.ToJustFloat64, axis: AnyShape | None = None) -> np.float64: ...
+@overload  # 0d ~complex128
+def softmax(x: onp.ToJustComplex128, axis: AnyShape | None = None) -> np.complex128: ...
+@overload  # nd T@inexact
+def softmax(x: onp.ToArrayND[_InexactT, _InexactT], axis: AnyShape | None = None) -> onp.ArrayND[_InexactT]: ...
+@overload  # nd +float64
+def softmax(x: onp.ToIntND | onp.ToJustFloat64_ND, axis: AnyShape | None = None) -> onp.ArrayND[np.float64]: ...
+@overload  # nd ~complex128
+def softmax(x: onp.ToJustComplex128_ND, axis: AnyShape | None = None) -> onp.ArrayND[np.complex128]: ...
+@overload  # 0d float fallback
+def softmax(x: onp.ToFloat, axis: AnyShape | None = None) -> np.float64 | Any: ...
+@overload  # 0d complex fallback
+def softmax(x: onp.ToComplex, axis: AnyShape | None = None) -> np.complex128 | Any: ...
+@overload  # nd float fallback
+def softmax(x: onp.ToFloatND, axis: AnyShape | None = None) -> onp.ArrayND[np.float64 | Any]: ...
+@overload  # nd complex fallback
+def softmax(x: onp.ToComplexND, axis: AnyShape | None = None) -> onp.ArrayND[np.complex128 | Any]: ...
 
-#
-@overload
-def log_softmax(x: onp.ToFloat, axis: AnyShape | None = None) -> np.float64: ...
-@overload
-def log_softmax(x: onp.ToFloatND, axis: AnyShape | None = None) -> onp.ArrayND[np.float64]: ...
-@overload
-def log_softmax(x: onp.ToComplex, axis: AnyShape | None = None) -> np.float64 | np.complex128: ...
-@overload
-def log_softmax(x: onp.ToComplexND, axis: AnyShape | None = None) -> onp.ArrayND[np.float64 | np.complex128]: ...
+# NOTE: keep in sync with `softmax`
+@overload  # T
+def log_softmax(x: _InexactOrArrayT, axis: AnyShape | None = None) -> _InexactOrArrayT: ...  # type: ignore[overload-overlap]
+@overload  # 0d +float64
+def log_softmax(x: onp.ToInt | onp.ToJustFloat64, axis: AnyShape | None = None) -> np.float64: ...
+@overload  # 0d ~complex128
+def log_softmax(x: onp.ToJustComplex128, axis: AnyShape | None = None) -> np.complex128: ...
+@overload  # nd T@inexact
+def log_softmax(x: onp.ToArrayND[_InexactT, _InexactT], axis: AnyShape | None = None) -> onp.ArrayND[_InexactT]: ...
+@overload  # nd +float64
+def log_softmax(x: onp.ToIntND | onp.ToJustFloat64_ND, axis: AnyShape | None = None) -> onp.ArrayND[np.float64]: ...
+@overload  # nd ~complex128
+def log_softmax(x: onp.ToJustComplex128_ND, axis: AnyShape | None = None) -> onp.ArrayND[np.complex128]: ...
+@overload  # 0d float fallback
+def log_softmax(x: onp.ToFloat, axis: AnyShape | None = None) -> np.float64 | Any: ...
+@overload  # 0d complex fallback
+def log_softmax(x: onp.ToComplex, axis: AnyShape | None = None) -> np.complex128 | Any: ...
+@overload  # nd float fallback
+def log_softmax(x: onp.ToFloatND, axis: AnyShape | None = None) -> onp.ArrayND[np.float64 | Any]: ...
+@overload  # nd complex fallback
+def log_softmax(x: onp.ToComplexND, axis: AnyShape | None = None) -> onp.ArrayND[np.complex128 | Any]: ...

tests/special/test_logsumexp.pyi

@@ -3,7 +3,7 @@ from typing import Any, assert_type
 import numpy as np
 import numpy.typing as npt
 
-from scipy.special import logsumexp
+from scipy.special import logsumexp, softmax
 
 py_f_0d: float
 py_c_0d: complex
@@ -75,3 +75,16 @@ assert_type(
 assert_type(
     logsumexp(c64_1d, axis=0, return_sign=True), tuple[npt.NDArray[np.floating[Any]] | Any, npt.NDArray[np.complex64] | Any]
 )
+
+###
+# softmax (equiv log_softmax)
+
+assert_type(softmax(py_f_0d), np.float64)
+assert_type(softmax(py_c_0d), np.complex128)
+assert_type(softmax(f16_0d), np.float16)
+assert_type(softmax(c64_0d), np.complex64)
+
+assert_type(softmax(py_f_1d), npt.NDArray[np.float64])
+assert_type(softmax(py_c_1d), npt.NDArray[np.complex128])
+assert_type(softmax(f16_1d), np.ndarray[tuple[int], np.dtype[np.float16]])
+assert_type(softmax(c64_1d), np.ndarray[tuple[int], np.dtype[np.complex64]])