✨ linalg: improved eigh return type inference for float64 (#755)

Author: jorenham

scipy-stubs/linalg/_decomp.pyi

@@ -384,7 +384,22 @@ def eig(
     ...
 
 #
-@overload  # float, eigvals_only: False = ...
+@overload  # +float64, eigvals_only: False = ...
+def eigh(  #
+    a: onp.ToArrayND[float, np.float64 | np.longdouble | npc.integer64 | npc.integer32],
+    b: onp.ToFloat64_ND | None = None,
+    *,
+    lower: op.CanBool = True,
+    eigvals_only: onp.ToFalse = False,
+    overwrite_a: op.CanBool = False,
+    overwrite_b: op.CanBool = False,
+    type: _EigHType = 1,
+    check_finite: op.CanBool = True,
+    subset_by_index: _EigHSubsetByIndex | None = None,
+    subset_by_value: _EigHSubsetByValue | None = None,
+    driver: _DriverEV | _DriverGV | None = None,
+) -> tuple[onp.ArrayND[np.float64], onp.ArrayND[np.float64]]: ...
+@overload  # +float, eigvals_only: False = ...
 def eigh(
     a: onp.ToFloatND,
     b: onp.ToFloatND | None = None,
@@ -399,22 +414,22 @@ def eigh(
     subset_by_value: _EigHSubsetByValue | None = None,
     driver: _DriverEV | _DriverGV | None = None,
 ) -> tuple[_FloatND, _FloatND]: ...
-@overload  # float, eigvals_only: True
+@overload  # ~complex, eigvals_only: False = ...
 def eigh(
-    a: onp.ToFloatND,
-    b: onp.ToFloatND | None = None,
+    a: onp.ToJustComplexND,
+    b: onp.ToComplexND | None = None,
     *,
     lower: op.CanBool = True,
-    eigvals_only: onp.ToTrue,
+    eigvals_only: onp.ToFalse = False,
     overwrite_a: op.CanBool = False,
     overwrite_b: op.CanBool = False,
     type: _EigHType = 1,
     check_finite: op.CanBool = True,
     subset_by_index: _EigHSubsetByIndex | None = None,
     subset_by_value: _EigHSubsetByValue | None = None,
-    driver: _DriverEV | _EigHSubsetByValue | None = None,
-) -> _FloatND: ...
-@overload  # complex, eigvals_only: False = ...
+    driver: _DriverEV | _DriverGV | None = None,
+) -> tuple[_FloatND, _ComplexND]: ...
+@overload  # +complex, eigvals_only: False = ...
 def eigh(
     a: onp.ToComplexND,
     b: onp.ToComplexND | None = None,
@@ -429,10 +444,10 @@ def eigh(
     subset_by_value: _EigHSubsetByValue | None = None,
     driver: _DriverEV | _DriverGV | None = None,
 ) -> tuple[_FloatND, _InexactND]: ...
-@overload  # complex, eigvals_only: True
+@overload  # +complex128, eigvals_only: True
 def eigh(
-    a: onp.ToComplexND,
-    b: onp.ToComplexND | None = None,
+    a: onp.ToArrayND[float, npc.inexact80 | npc.number64 | npc.integer32],
+    b: onp.ToComplex128_ND | None = None,
     *,
     lower: op.CanBool = True,
     eigvals_only: onp.ToTrue,
@@ -443,22 +458,22 @@ def eigh(
     subset_by_index: _EigHSubsetByIndex | None = None,
     subset_by_value: _EigHSubsetByValue | None = None,
     driver: _DriverEV | _EigHSubsetByValue | None = None,
-) -> _FloatND: ...
-@overload  # complex, eigvals_only: CanBool (catch-all)
+) -> onp.ArrayND[np.float64]: ...
+@overload  # +complex, eigvals_only: True
 def eigh(
     a: onp.ToComplexND,
     b: onp.ToComplexND | None = None,
     *,
-    lower: op.CanBool,
-    eigvals_only: op.CanBool,
+    lower: op.CanBool = True,
+    eigvals_only: onp.ToTrue,
     overwrite_a: op.CanBool = False,
     overwrite_b: op.CanBool = False,
     type: _EigHType = 1,
     check_finite: op.CanBool = True,
     subset_by_index: _EigHSubsetByIndex | None = None,
     subset_by_value: _EigHSubsetByValue | None = None,
     driver: _DriverEV | _EigHSubsetByValue | None = None,
-) -> _FloatND | tuple[_FloatND, _InexactND]: ...
+) -> _FloatND: ...
 
 #
 @overload  # float, eigvals_only: False = ..., select: _SelectA = ...

scipy-stubs/linalg/_matfuncs.pyi

@@ -41,6 +41,8 @@ _Complex128ND: TypeAlias = onp.ArrayND[np.complex128]
 _ComplexND: TypeAlias = onp.ArrayND[npc.complexfloating]
 _InexactND: TypeAlias = onp.ArrayND[npc.inexact]
 
+_AsFloat64ND: TypeAlias = onp.ToArrayND[float, np.float64 | npc.integer | np.bool_]
+
 ###
 
 eps: Final[np.float64] = ...  # undocumented
@@ -77,7 +79,7 @@ def fractional_matrix_power(A: onp.ToComplexND, t: onp.ToFloat) -> onp.ArrayND[A
 
 #
 @overload
-def sqrtm(A: onp.ToIntND | onp.ToJustFloat64_ND) -> _Float64ND: ...
+def sqrtm(A: _AsFloat64ND) -> _Float64ND: ...
 @overload
 def sqrtm(A: onp.ToFloatND) -> _FloatND: ...
 @overload
@@ -106,7 +108,7 @@ def sqrtm(A: onp.ToComplexND, disp: onp.ToFalse, blocksize: int) -> tuple[_Inexa
 
 # NOTE: return dtype depends on the sign of the values
 @overload  # +integer | ~float64
-def logm(A: onp.ToIntND | onp.ToJustFloat64_ND) -> _Float64ND | _Complex128ND: ...
+def logm(A: _AsFloat64ND) -> _Float64ND | _Complex128ND: ...
 @overload  # +floating
 def logm(A: onp.ToFloatND) -> _InexactND: ...
 @overload  # ~complex128
@@ -127,7 +129,7 @@ def logm(A: onp.ToComplexND, disp: onp.ToFalse) -> tuple[_InexactND, float]: ...
 
 #
 @overload  # +integer | ~float64
-def expm(A: onp.ToIntND | onp.ToJustFloat64_ND) -> _Float64ND: ...
+def expm(A: _AsFloat64ND) -> _Float64ND: ...
 @overload  # +floating
 def expm(A: onp.ToFloatND) -> _FloatND: ...
 @overload  # ~complex128
@@ -142,7 +144,7 @@ def _exp_sinch(x: onp.ArrayND[_ComplexT, _ShapeT]) -> onp.ArrayND[_ComplexT, _Sh
 
 #
 @overload  # +integer | ~float64
-def cosm(A: onp.ToIntND | onp.ToJustFloat64_ND) -> _Float64ND: ...
+def cosm(A: _AsFloat64ND) -> _Float64ND: ...
 @overload  # +floating
 def cosm(A: onp.ToFloatND) -> _FloatND: ...
 @overload  # ~complex128
@@ -154,7 +156,7 @@ def cosm(A: onp.ToComplexND) -> _InexactND: ...
 
 #
 @overload  # +integer | ~float64
-def sinm(A: onp.ToIntND | onp.ToJustFloat64_ND) -> _Float64ND: ...
+def sinm(A: _AsFloat64ND) -> _Float64ND: ...
 @overload  # +floating
 def sinm(A: onp.ToFloatND) -> _FloatND: ...
 @overload  # ~complex128
@@ -166,7 +168,7 @@ def sinm(A: onp.ToComplexND) -> _InexactND: ...
 
 #
 @overload  # +integer | ~float64
-def tanm(A: onp.ToIntND | onp.ToJustFloat64_ND) -> _Float64ND: ...
+def tanm(A: _AsFloat64ND) -> _Float64ND: ...
 @overload  # +floating
 def tanm(A: onp.ToFloatND) -> _FloatND: ...
 @overload  # ~complex128
@@ -178,7 +180,7 @@ def tanm(A: onp.ToComplexND) -> _InexactND: ...
 
 #
 @overload  # +integer | ~float64
-def coshm(A: onp.ToIntND | onp.ToJustFloat64_ND) -> _Float64ND: ...
+def coshm(A: _AsFloat64ND) -> _Float64ND: ...
 @overload  # +floating
 def coshm(A: onp.ToFloatND) -> _FloatND: ...
 @overload  # ~complex128
@@ -190,7 +192,7 @@ def coshm(A: onp.ToComplexND) -> _InexactND: ...
 
 #
 @overload  # +integer | ~float64
-def sinhm(A: onp.ToIntND | onp.ToJustFloat64_ND) -> _Float64ND: ...
+def sinhm(A: _AsFloat64ND) -> _Float64ND: ...
 @overload  # +floating
 def sinhm(A: onp.ToFloatND) -> _FloatND: ...
 @overload  # ~complex128
@@ -202,7 +204,7 @@ def sinhm(A: onp.ToComplexND) -> _InexactND: ...
 
 #
 @overload  # +integer | ~float64
-def tanhm(A: onp.ToIntND | onp.ToJustFloat64_ND) -> _Float64ND: ...
+def tanhm(A: _AsFloat64ND) -> _Float64ND: ...
 @overload  # +floating
 def tanhm(A: onp.ToFloatND) -> _FloatND: ...
 @overload  # ~complex128
@@ -228,7 +230,7 @@ def funm(A: onp.CanArrayND[_InexactT], func: _FuncND[_InexactT], disp: onp.ToFal
 
 #
 @overload  # +float64
-def signm(A: onp.ToIntND | onp.ToFloat64_ND) -> _Float64ND: ...
+def signm(A: _AsFloat64ND) -> _Float64ND: ...
 @overload  # +floating
 def signm(A: onp.ToFloatND) -> _FloatND: ...
 @overload  # +complexfloating
@@ -248,7 +250,7 @@ def signm(A: onp.ToComplexND, disp: onp.ToTrue) -> _InexactND: ...
 def signm(A: onp.ToComplexND, disp: onp.ToFalse) -> tuple[_InexactND, np.float64]: ...
 
 #
-@overload  # +integer, +integer
+@overload  # +integer | ~float64
 def khatri_rao(a: onp.ToIntND, b: onp.ToIntND) -> _IntND: ...
 @overload  # +float64, ~float64
 def khatri_rao(a: onp.ToFloat64_ND, b: onp.ToJustFloat64_ND) -> _Float64ND: ...