🐛stats.unitary_group.rvs: fix return dtype

scipy-stubs/stats/_multivariate.pyi

@@ -31,6 +31,7 @@ __all__ = [
 ]
 
 _ScalarT = TypeVar("_ScalarT", bound=np.generic, default=np.float64)
+_ScalarT_co = TypeVar("_ScalarT_co", bound=np.generic, default=np.float64, covariant=True)
 
 # TODO(@jorenham): rename as {}T_co
 _RVG_co = TypeVar("_RVG_co", bound=multi_rv_generic, default=multi_rv_generic, covariant=True)
@@ -363,24 +364,26 @@ class multinomial_frozen(multi_rv_frozen[multinomial_gen]):
     def rvs(self, /, size: onp.AtLeast1D | int = 1, random_state: onp.random.ToRNG | None = None) -> _Array2ND: ...
 
 @type_check_only
-class _group_rv_gen_mixin(Generic[_RVF_co]):
+class _group_rv_gen_mixin(Generic[_RVF_co, _ScalarT_co]):
     def __call__(self, /, dim: onp.ToJustInt | None = None, seed: onp.random.ToRNG | None = None) -> _RVF_co: ...
     def rvs(
         self, /, dim: onp.ToJustInt, size: onp.ToJustInt | None = 1, random_state: onp.random.ToRNG | None = None
-    ) -> onp.Array3D[np.float64]: ...
+    ) -> onp.Array3D[_ScalarT_co]: ...
 
 @type_check_only
-class _group_rv_frozen_mixin:
+class _group_rv_frozen_mixin(Generic[_ScalarT_co]):
     dim: onp.ToJustInt
     def __init__(self, /, dim: onp.ToJustInt | None = None, seed: onp.random.ToRNG | None = None) -> None: ...
-    def rvs(self, /, size: onp.ToJustInt | None = 1, random_state: onp.random.ToRNG | None = None) -> onp.Array3D[np.float64]: ...
+    def rvs(
+        self, /, size: onp.ToJustInt | None = 1, random_state: onp.random.ToRNG | None = None
+    ) -> onp.Array3D[_ScalarT_co]: ...
 
 class special_ortho_group_gen(_group_rv_gen_mixin[special_ortho_group_frozen], multi_rv_generic): ...
 class special_ortho_group_frozen(_group_rv_frozen_mixin, multi_rv_frozen[special_ortho_group_gen]): ...
 class ortho_group_gen(_group_rv_gen_mixin[ortho_group_frozen], multi_rv_generic): ...
 class ortho_group_frozen(_group_rv_frozen_mixin, multi_rv_frozen[ortho_group_gen]): ...
-class unitary_group_gen(_group_rv_gen_mixin[unitary_group_frozen], multi_rv_generic): ...
-class unitary_group_frozen(_group_rv_frozen_mixin, multi_rv_frozen[unitary_group_gen]): ...
+class unitary_group_gen(_group_rv_gen_mixin[unitary_group_frozen, np.complex128], multi_rv_generic): ...
+class unitary_group_frozen(_group_rv_frozen_mixin[np.complex128], multi_rv_frozen[unitary_group_gen]): ...
 
 #
 class uniform_direction_gen(multi_rv_generic):

tests/stats/test_multivariate.pyi

@@ -0,0 +1,77 @@
+from typing import assert_type
+
+import numpy as np
+
+from scipy.stats import (
+    dirichlet,
+    invwishart,
+    matrix_normal,
+    multinomial,
+    multivariate_hypergeom,
+    multivariate_normal,
+    multivariate_t,
+    normal_inverse_gamma,
+    ortho_group,
+    random_correlation,
+    random_table,
+    special_ortho_group,
+    uniform_direction,
+    unitary_group,
+    vonmises_fisher,
+    wishart,
+)
+
+###
+# rvs dtype checks
+
+# the pyright ignore comments are needed due to  https://github.com/microsoft/pyright/issues/11127
+
+assert_type(multivariate_normal.rvs().dtype, np.dtype[np.float64])
+assert_type(multivariate_normal().rvs().dtype, np.dtype[np.float64])
+
+assert_type(matrix_normal.rvs().dtype, np.dtype[np.float64])
+assert_type(matrix_normal().rvs().dtype, np.dtype[np.float64])
+
+assert_type(dirichlet.rvs([1, 2]).dtype, np.dtype[np.float64])  # pyright: ignore[reportUnknownMemberType]
+assert_type(dirichlet([1, 2]).rvs().dtype, np.dtype[np.float64])  # pyright: ignore[reportUnknownMemberType]
+
+assert_type(wishart.rvs(1, 1).dtype, np.dtype[np.float64])
+assert_type(wishart().rvs().dtype, np.dtype[np.float64])
+
+assert_type(invwishart.rvs(1, 1).dtype, np.dtype[np.float64])
+assert_type(invwishart().rvs().dtype, np.dtype[np.float64])
+
+assert_type(multinomial.rvs([1], [0.5]).dtype, np.dtype[np.float64])  # pyright: ignore[reportUnknownMemberType]
+assert_type(multinomial([1], [0.5]).rvs().dtype, np.dtype[np.float64])  # pyright: ignore[reportUnknownMemberType]
+
+assert_type(ortho_group.rvs(3).dtype, np.dtype[np.float64])
+assert_type(ortho_group().rvs(3).dtype, np.dtype[np.float64])
+
+assert_type(special_ortho_group.rvs(3).dtype, np.dtype[np.float64])
+assert_type(special_ortho_group().rvs(3).dtype, np.dtype[np.float64])
+
+assert_type(unitary_group.rvs(3).dtype, np.dtype[np.complex128])
+assert_type(unitary_group().rvs(3).dtype, np.dtype[np.complex128])
+
+assert_type(uniform_direction.rvs(3).dtype, np.dtype[np.float64])
+assert_type(uniform_direction(1).rvs(3).dtype, np.dtype[np.float64])
+
+assert_type(random_correlation.rvs([1, 1]).dtype, np.dtype[np.float64])
+assert_type(random_correlation([1, 1]).rvs().dtype, np.dtype[np.float64])
+
+assert_type(multivariate_t.rvs().dtype, np.dtype[np.float64])  # pyright: ignore[reportUnknownMemberType]
+assert_type(multivariate_t().rvs().dtype, np.dtype[np.float64])  # pyright: ignore[reportUnknownMemberType]
+
+assert_type(multivariate_hypergeom.rvs([1], 1).dtype, np.dtype[np.float64])  # pyright: ignore[reportUnknownMemberType]
+assert_type(multivariate_hypergeom([1], 1).rvs().dtype, np.dtype[np.float64])  # pyright: ignore[reportUnknownMemberType]
+
+assert_type(random_table.rvs([1, 2], [2, 1]).dtype, np.dtype[np.float64])
+assert_type(random_table([1, 2], [2, 1]).rvs().dtype, np.dtype[np.float64])
+
+# `dirichlet_multinomial` has no `rvs` method
+
+assert_type(vonmises_fisher.rvs([0.8, 0.6]).dtype, np.dtype[np.float64])  # pyright: ignore[reportUnknownMemberType]
+assert_type(vonmises_fisher([0.8, 0.6]).rvs().dtype, np.dtype[np.float64])  # pyright: ignore[reportUnknownMemberType]
+
+assert_type(normal_inverse_gamma.rvs()[0].dtype, np.dtype[np.float64])
+assert_type(normal_inverse_gamma().rvs()[0].dtype, np.dtype[np.float64])