From 48e17d40f8304bf336c8c709183fb88dea9b59f5 Mon Sep 17 00:00:00 2001 From: AlessandroRuggiero Date: Wed, 11 Nov 2026 21:32:47 +0100 Subject: [PATCH] Infer NumPy Generator.normal result shapes --- .../numpy-stubs/random/__init__.pyi | 96 ++++++++++++++++++- .../pyrefly-numpy-stubs/test/test_random.py | 59 ++++++++++++ 2 files changed, 152 insertions(+), 3 deletions(-) diff --git a/tensor-shapes/pyrefly-numpy-stubs/numpy-stubs/random/__init__.pyi b/tensor-shapes/pyrefly-numpy-stubs/numpy-stubs/random/__init__.pyi index 701d963f87..c3775cd5f3 100644 --- a/tensor-shapes/pyrefly-numpy-stubs/numpy-stubs/random/__init__.pyi +++ b/tensor-shapes/pyrefly-numpy-stubs/numpy-stubs/random/__init__.pyi @@ -3,13 +3,19 @@ # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. -from typing import Literal, overload +from typing import Any, Literal, overload from numpy._shapes import ( multivariate_normal_scalar_size_shape, multivariate_normal_shape, ) -from numpy._typing import _ArrayLikeInt_co, _ShapeLike, ArrayLike +from numpy._typing import ( + _ArrayLikeFloat_co, + _ArrayLikeInt_co, + _FloatLike_co, + _ShapeLike, + ArrayLike, +) from numpy.random._generator import ( _ArrayF32, _ArrayF64, @@ -17,6 +23,7 @@ from numpy.random._generator import ( _DTypeLikeF64, _DTypeLikeFloat, _MethodExp, + _NDArrayLikeFloat, Generator as _NumpyGenerator, ) @@ -90,7 +97,7 @@ from numpy.random.mtrand import ( weibull as weibull, zipf as zipf, ) -from shape_extensions import Int, IntTuple, IntVar +from shape_extensions import broadcast, Int, IntTuple, IntVar from .. import dtype, float32, float64, ndarray @@ -336,6 +343,89 @@ class Generator(_NumpyGenerator): *, out: ArrayT, ) -> ArrayT: ... + @overload + def normal( + self, /, loc: _FloatLike_co = 0.0, scale: _FloatLike_co = 1.0, size: None = None + ) -> float: ... + # With `size`, the result has exactly that shape, and NumPy raises if `loc` or + # `scale` does not broadcast to it. + @overload + def normal[N: IntVar]( + self, /, loc: _ArrayLikeFloat_co, scale: _ArrayLikeFloat_co, size: Int[N] + ) -> ndarray[[N], dtype[float64]]: ... + @overload + def normal[Shape: IntTuple]( + self, /, loc: _ArrayLikeFloat_co, scale: _ArrayLikeFloat_co, size: Shape + ) -> ndarray[Shape, dtype[float64]]: ... + @overload + def normal[N: IntVar]( + self, + /, + loc: _ArrayLikeFloat_co = 0.0, + scale: _ArrayLikeFloat_co = 1.0, + *, + size: Int[N], + ) -> ndarray[[N], dtype[float64]]: ... + @overload + def normal[Shape: IntTuple]( + self, + /, + loc: _ArrayLikeFloat_co = 0.0, + scale: _ArrayLikeFloat_co = 1.0, + *, + size: Shape, + ) -> ndarray[Shape, dtype[float64]]: ... + # Without `size`, the result has the broadcast shape of `loc` and `scale`. + @overload + def normal[LocShape: IntTuple = [], ScaleShape: IntTuple = []]( + self, + /, + loc: ndarray[LocShape] | float = 0.0, + scale: ndarray[ScaleShape] | float = 1.0, + size: None = None, + ) -> ndarray[broadcast(LocShape, ScaleShape), dtype[float64]]: ... + @overload + def normal( + self, /, loc: _ArrayLikeFloat_co, scale: _ArrayLikeFloat_co, size: _ShapeLike + ) -> _ArrayF64: ... + @overload + def normal( + self, /, loc: _ArrayLikeFloat_co, scale: _NDArrayLikeFloat, size: None + ) -> _ArrayF64: ... + @overload + def normal( + self, + /, + loc: _ArrayLikeFloat_co = 0.0, + scale: _ArrayLikeFloat_co = 1.0, + *, + size: _ShapeLike, + ) -> _ArrayF64: ... + @overload + def normal( + self, + /, + loc: _ArrayLikeFloat_co = 0.0, + *, + scale: _NDArrayLikeFloat, + size: None = None, + ) -> _ArrayF64: ... + @overload + def normal( + self, + /, + loc: _NDArrayLikeFloat, + scale: _ArrayLikeFloat_co = 1.0, + size: None = None, + ) -> _ArrayF64: ... + @overload + def normal( + self, + /, + loc: _ArrayLikeFloat_co = 0.0, + scale: _ArrayLikeFloat_co = 1.0, + size: None = None, + ) -> _ArrayF64 | Any: ... def default_rng( seed: _ArrayLikeInt_co diff --git a/tensor-shapes/pyrefly-numpy-stubs/test/test_random.py b/tensor-shapes/pyrefly-numpy-stubs/test/test_random.py index 193b1bfdf7..64d07b59dc 100644 --- a/tensor-shapes/pyrefly-numpy-stubs/test/test_random.py +++ b/tensor-shapes/pyrefly-numpy-stubs/test/test_random.py @@ -94,6 +94,65 @@ def check_standard_normal_and_standard_exponential_symbolic_size[N: IntVar, M: I ) +def check_normal_returns_float_without_size() -> None: + rng = np.random.default_rng(0) + assert_type(rng.normal(), float) + assert_type(rng.normal(1.0, 2.0), float) + + +def test_normal_int_size() -> None: + rng = np.random.default_rng(0) + result = rng.normal(size=5) + assert_type(result, np.ndarray[[5], np.dtype[np.float64]]) + assert_shape(result.shape, (5,)) + assert_shape(rng.normal(0.0, 1.0, 5).shape, (5,)) + + +def test_normal_tuple_size() -> None: + rng = np.random.default_rng(0) + result = rng.normal(0.0, 1.0, (2, 3)) + assert_type(result, np.ndarray[[2, 3], np.dtype[np.float64]]) + assert_shape(result.shape, (2, 3)) + assert_shape(rng.normal(size=(2, 3)).shape, (2, 3)) + + +def test_normal_array_parameters_with_size() -> None: + rng = np.random.default_rng(0) + result = rng.normal(np.zeros(3), 1.0, size=(2, 3)) + assert_type(result, np.ndarray[[2, 3], np.dtype[np.float64]]) + assert_shape(result.shape, (2, 3)) + assert_shape(rng.normal(np.zeros((10, 1)), size=(10, 5)).shape, (10, 5)) + + +def test_normal_array_parameters_without_size() -> None: + rng = np.random.default_rng(0) + result = rng.normal(np.zeros(3), np.ones((2, 1))) + assert_type(result, np.ndarray[[2, 3], np.dtype[np.float64]]) + assert_shape(result.shape, (2, 3)) + assert_shape(rng.normal(np.zeros(3)).shape, (3,)) + assert_shape(rng.normal(scale=np.ones((2, 1))).shape, (2, 1)) + + +def test_normal_list_size_and_parameters() -> None: + rng = np.random.default_rng(0) + assert_shape(rng.normal(size=[2, 3]).shape, IntTuple, runtime=(2, 3)) + assert_shape(rng.normal([0.0, 1.0, 2.0]).shape, IntTuple, runtime=(3,)) + + +def test_normal_mismatched_parameters() -> None: + rng = np.random.default_rng(0) + with assert_raises(ValueError): + rng.normal(np.zeros(3), np.ones(2)) # E: Cannot broadcast + + +def check_normal_symbolic_shapes[N: IntVar, M: IntVar]( + n: Int[N], loc: np.ndarray[[N, 1]], scale: np.ndarray[[M]] +) -> None: + rng = np.random.default_rng(0) + assert_type(rng.normal(size=n), np.ndarray[[N], np.dtype[np.float64]]) + assert_type(rng.normal(loc, scale), np.ndarray[[N, M], np.dtype[np.float64]]) + + def test_multivariate_normal_without_size() -> None: rng = np.random.default_rng(0) result = rng.multivariate_normal(np.zeros(3), np.eye(3))