Add generic shape typing for multivariate frozen distributions in stats._multivariate by Aniketsy · Pull Request #1549 · scipy/scipy-stubs

Aniketsy · 2026-04-14T10:10:01Z

Fixes #406

github-actions · 2026-04-14T10:16:39Z

`mypy_primer` results

✅ No ecosystem changes detected

Aniketsy · 2026-04-14T11:17:05Z

ahhh 😭 , I did run locally

aniket@DESKTOP-074O80J:/mnt/d/scipy-stubs/scipy-stubs$ uv run mypy scipy-stubs
Success: no issues found in 589 source files
aniket@DESKTOP-074O80J:/mnt/d/scipy-stubs/scipy-stubs$ uv run basedpyright
0 errors, 0 warnings, 0 notes

i'll fix these ci fails-- shortly.

jorenham · 2026-04-14T11:21:51Z

uvx tox p might help prevent some CI failures

Aniketsy · 2026-04-18T12:02:06Z

@jorenham could you please review this, when you get chance and share some pointers.
Also i have one small suggestion, can we implement autofix, pre-commit for ruff, as previously i run ruff check it passed locally, but here i got ci-failures for ruff. (not sure if we should implement 🤷 )

I took help from llm and as per suggestion added comments (# pyrefly: ignore [inconsistent-inheritance]), please let me know if these should be reverted as we're getting ci-failures. Thanks!

jorenham · 2026-04-18T12:09:00Z

@jorenham could you please review this, when you get chance and share some pointers.

I'll have a look soon

Also i have one small suggestion, can we implement autofix, pre-commit for ruff, as previously i run ruff check it passed locally, but here i got ci-failures for ruff. (not sure if we should implement 🤷 )

That's exactly what we use lefthook for :)
https://github.com/scipy/scipy-stubs/blob/master/CONTRIBUTING.md#lefthook

jorenham

Hmm, this isn't quite what I had in mind. Apologies for not explaining it better in the issue.

It's probably easiest to close this PR and start over, separating the issue you found in #406 (comment), and the shape-typing improvements. And it's probably easiest to do implement the shape-typing improvements one distribution at a time.
How about I open a PR myself for one of those distributions so you understand what I had mind, and leave the rest for you?

jorenham · 2026-04-18T14:35:32Z

+_K = TypeVar("_K", bound=int, default=int)
+_M = TypeVar("_M", bound=int, default=int)
+_N = TypeVar("_N", bound=int, default=int)
+_P = TypeVar("_P", bound=int, default=int)
+_D = TypeVar("_D", bound=int, default=int)


NumPy does not support using subtypes of int in shape-types, so let's just stick with plain int as well.

jorenham · 2026-04-18T14:40:27Z

-# TODO(@jorenham): Generic shape-type for mean and cov, so that we can determine whether the methods return scalars or arrays.
-# https://github.com/scipy/scipy-stubs/issues/406
-class multivariate_normal_frozen(multi_rv_frozen[multivariate_normal_gen]):
+class multivariate_normal_frozen(multi_rv_frozen[multivariate_normal_gen], Generic[_K]):


Hmm I just checked, and apparently multivariate_normal only supports 1d mean and 2d covariance, unlike the issue suggest. So there's no reason to make this generic after all.

jorenham · 2026-04-18T14:41:41Z

+        rowcov: onp.ToFloat | onp.ToFloat1D | onp.ToFloat2D = 1,
+        colcov: onp.ToFloat | onp.ToFloat1D | onp.ToFloat2D = 1,


Let's separate this fix and the shape-typing improvements. It'd otherwise be confusing if a PR is both a fix and a feature, and would fit in the changelog format.

jorenham · 2026-04-18T14:46:41Z

+    def entropy(self, /, rowcov: onp.ToFloat | onp.ToFloat1D | onp.ToFloat2D = 1, colcov: onp.ToFloat | onp.ToFloat1D | onp.ToFloat2D = 1) -> np.float64: ...

-class matrix_normal_frozen(multi_rv_frozen[matrix_normal_gen]):
+class matrix_normal_frozen(multi_rv_frozen[matrix_normal_gen], Generic[_M, _N]):


Well, apparently I was wrong here too. Maybe something changed, but with the latest scipy it looks like matrix_normal only accepts 2d mean, rowcov and colcov. So sorry for the confusion, but here too is there no need to make the shape-type generic.

jorenham · 2026-04-18T14:51:47Z

 #
-class multinomial_frozen(multi_rv_frozen[multinomial_gen]):
-    def __init__(self, /, n: onp.ToJustIntND, p: onp.ToJustFloatND, seed: onp.random.ToRNG | None = None) -> None: ...
+class multinomial_frozen(multi_rv_frozen[multinomial_gen], Generic[_K]):


Ok so this should indeed be generic on its shape-type:

In [27]: multinomial(8, [0.3, 0.2, 0.5]).mean() Out[27]: array([2.4, 1.6, 4. ]) In [28]: multinomial([8,9], [0.3, 0.2, 0.5]).mean() Out[28]: array([[2.4, 1.6, 4. ], [2.7, 1.8, 4.5]]) In [29]: multinomial([[8,9],[6,5]], [0.3, 0.2, 0.5]).mean() Out[29]: array([[[2.4, 1.6, 4. ], [2.7, 1.8, 4.5]], [[1.8, 1.2, 3. ], [1.5, 1. , 2.5]]])

However, the generic type parameter should be

_ShapeT_co = TypeVar("`_ShapeT_co", bound=tuple[int, ...], default=tuple[Any, ...], covariant=True)

matching the shape-type of .mean()

jorenham · 2026-04-18T14:52:57Z


 @type_check_only
 class _group_rv_frozen_mixin(Generic[_ScalarT_co]):
-    __class_getitem__: ClassVar[None] = None


why did you remove this?

as it sets to None -disables generic , please correct me if i misunderstood here.

jorenham · 2026-04-18T14:53:33Z

 @type_check_only
 class _group_rv_frozen_mixin(Generic[_ScalarT_co]):
-    __class_getitem__: ClassVar[None] = None
-


please avoid unnecessary formatting changes like this one and the ones below

Aniketsy · 2026-04-18T15:57:48Z

How about I open a PR myself for one of those distributions so you understand what I had mind, and leave the rest for you?

sure, please go for it, and then i can follow you ..

It's probably easiest to close this PR and start over, separating the issue you found in #406 (comment),

okay , then i'll close this one and open a separate PR for the issue i found and then follow up after your PR.

Thanks! for giving your time in review and apologies for creating extra work .

jorenham · 2026-04-18T16:00:52Z

sure, please go for it, and then i can follow you ..

I'll get myself something to eat and then start working on it 👌

Aniketsy · 2026-04-18T16:08:14Z

I'll get myself something to eat and then start working on it 👌

sure 😊, please ping me also once its done , I'd love to see and understand . thanks ! .... closing this PR

add generic shape types for multivariate frozen distributions

19c602c

Aniketsy marked this pull request as draft April 14, 2026 10:15

jorenham added scipy.stats is: feature labels Apr 14, 2026

Aniketsy added 2 commits April 14, 2026 10:26

add generic shape types for multivariate frozen distributions

0da08db

add generic shape types for multivariate frozen distributions

40502e2

Aniketsy added 4 commits April 15, 2026 17:18

add generic shape types for multivariate frozen distributions

dfa3a0c

revert __class_getitem__

b6cdc14

add generic shape types for multivariate frozen distributions

194981a

remove disable generic and add comments

44641bf

Aniketsy marked this pull request as ready for review April 18, 2026 11:47

jorenham requested changes Apr 18, 2026

View reviewed changes

Aniketsy closed this Apr 18, 2026

jorenham mentioned this pull request Apr 18, 2026

✨ stats: multinomial shape-typing #1566

Merged

		rowcov: onp.ToFloat \| onp.ToFloat1D \| onp.ToFloat2D = 1,
		colcov: onp.ToFloat \| onp.ToFloat1D \| onp.ToFloat2D = 1,

Uh oh!

Conversation

Aniketsy commented Apr 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions Bot commented Apr 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

mypy_primer results

Uh oh!

Aniketsy commented Apr 14, 2026

Uh oh!

jorenham commented Apr 14, 2026

Uh oh!

Aniketsy commented Apr 18, 2026

Uh oh!

jorenham commented Apr 18, 2026

Uh oh!

jorenham left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Aniketsy commented Apr 18, 2026

Uh oh!

jorenham commented Apr 18, 2026

Uh oh!

Aniketsy commented Apr 18, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Aniketsy commented Apr 14, 2026 •

edited

Loading

github-actions Bot commented Apr 14, 2026 •

edited

Loading

`mypy_primer` results