From dae23243f9ecec21720adc877db088e21cc6d323 Mon Sep 17 00:00:00 2001 From: AlessandroRuggiero Date: Sat, 14 Nov 2026 15:30:30 +0100 Subject: [PATCH] Infer NumPy Generator.integers result shapes --- .../numpy-stubs/random/__init__.pyi | 157 +++++++++++++++++- .../pyrefly-numpy-stubs/test/test_random.py | 80 +++++++++ 2 files changed, 236 insertions(+), 1 deletion(-) 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 c3775cd5f3..5f5ec3e416 100644 --- a/tensor-shapes/pyrefly-numpy-stubs/numpy-stubs/random/__init__.pyi +++ b/tensor-shapes/pyrefly-numpy-stubs/numpy-stubs/random/__init__.pyi @@ -12,9 +12,13 @@ from numpy._shapes import ( from numpy._typing import ( _ArrayLikeFloat_co, _ArrayLikeInt_co, + _DTypeLike, _FloatLike_co, + _Int64Codes, _ShapeLike, ArrayLike, + DTypeLike, + NDArray, ) from numpy.random._generator import ( _ArrayF32, @@ -22,6 +26,7 @@ from numpy.random._generator import ( _DTypeLikeF32, _DTypeLikeF64, _DTypeLikeFloat, + _DTypeLikeI64, _MethodExp, _NDArrayLikeFloat, Generator as _NumpyGenerator, @@ -99,7 +104,7 @@ from numpy.random.mtrand import ( ) from shape_extensions import broadcast, Int, IntTuple, IntVar -from .. import dtype, float32, float64, ndarray +from .. import bool as _bool, dtype, float32, float64, int64, integer, ndarray @overload def randn() -> float: ... @@ -166,6 +171,156 @@ class Generator(_NumpyGenerator): @overload def standard_cauchy(self, size: _ShapeLike) -> _ArrayF64: ... @overload + def integers[AnyIntT: (bool, int)]( + self, + low: int, + high: int | None = None, + size: None = None, + *, + dtype: type[AnyIntT], + endpoint: bool = False, + ) -> AnyIntT: ... + @overload + def integers[ScalarT: integer | _bool]( + self, + low: int, + high: int | None = None, + size: None = None, + *, + dtype: _DTypeLike[ScalarT], + endpoint: bool = False, + ) -> ScalarT: ... + @overload + def integers( + self, + low: int, + high: int | None = None, + size: None = None, + dtype: _DTypeLikeI64 = ..., + endpoint: bool = False, + ) -> int64: ... + @overload + def integers( + self, + low: int, + high: int | None = None, + size: None = None, + dtype: DTypeLike | None = ..., + endpoint: bool = False, + ) -> Any: ... + # With `size`, the result has exactly that shape, and NumPy raises if `low` or + # `high` does not broadcast to it. Without `dtype`, `ScalarT` defaults to `int64`. + @overload + def integers[N: IntVar, ScalarT: integer | _bool = int64]( + self, + low: _ArrayLikeInt_co, + high: _ArrayLikeInt_co | None, + size: Int[N], + dtype: _DTypeLike[ScalarT] | _Int64Codes = ..., + endpoint: bool = False, + ) -> ndarray[[N], dtype[ScalarT]]: ... + @overload + def integers[Shape: IntTuple, ScalarT: integer | _bool = int64]( + self, + low: _ArrayLikeInt_co, + high: _ArrayLikeInt_co | None, + size: Shape, + dtype: _DTypeLike[ScalarT] | _Int64Codes = ..., + endpoint: bool = False, + ) -> ndarray[Shape, dtype[ScalarT]]: ... + @overload + def integers[N: IntVar, ScalarT: integer | _bool = int64]( + self, + low: _ArrayLikeInt_co, + high: _ArrayLikeInt_co | None = None, + *, + size: Int[N], + dtype: _DTypeLike[ScalarT] | _Int64Codes = ..., + endpoint: bool = False, + ) -> ndarray[[N], dtype[ScalarT]]: ... + @overload + def integers[Shape: IntTuple, ScalarT: integer | _bool = int64]( + self, + low: _ArrayLikeInt_co, + high: _ArrayLikeInt_co | None = None, + *, + size: Shape, + dtype: _DTypeLike[ScalarT] | _Int64Codes = ..., + endpoint: bool = False, + ) -> ndarray[Shape, dtype[ScalarT]]: ... + # Without `size`, the result has the broadcast shape of `low` and `high`. + @overload + def integers[ + LowShape: IntTuple = [], + HighShape: IntTuple = [], + ScalarT: integer | _bool = int64, + ]( + self, + low: ndarray[LowShape] | int, + high: ndarray[HighShape] | int | None = None, + size: None = None, + dtype: _DTypeLike[ScalarT] | _Int64Codes = ..., + endpoint: bool = False, + ) -> ndarray[broadcast(LowShape, HighShape), dtype[ScalarT]]: ... + @overload + def integers[ScalarT: integer | _bool]( + self, + low: _ArrayLikeInt_co, + high: _ArrayLikeInt_co | None = None, + *, + size: _ShapeLike, + dtype: _DTypeLike[ScalarT], + endpoint: bool = False, + ) -> NDArray[ScalarT]: ... + @overload + def integers( + self, + low: _ArrayLikeInt_co, + high: _ArrayLikeInt_co | None = None, + *, + size: _ShapeLike, + dtype: _DTypeLikeI64 = ..., + endpoint: bool = False, + ) -> NDArray[int64]: ... + @overload + def integers( + self, + low: _ArrayLikeInt_co, + high: _ArrayLikeInt_co | None = None, + *, + size: _ShapeLike, + dtype: DTypeLike | None = ..., + endpoint: bool = False, + ) -> ndarray: ... + @overload + def integers[ScalarT: integer | _bool]( + self, + low: _ArrayLikeInt_co, + high: _ArrayLikeInt_co | None = None, + size: _ShapeLike | None = None, + *, + dtype: _DTypeLike[ScalarT], + endpoint: bool = False, + ) -> NDArray[ScalarT] | Any: ... + @overload + def integers( + self, + low: _ArrayLikeInt_co, + high: _ArrayLikeInt_co | None = None, + size: _ShapeLike | None = None, + dtype: _DTypeLikeI64 = ..., + endpoint: bool = False, + ) -> NDArray[int64] | Any: ... + @overload + def integers( + self, + low: _ArrayLikeInt_co, + high: _ArrayLikeInt_co | None = None, + size: _ShapeLike | None = None, + dtype: DTypeLike | None = ..., + endpoint: bool = False, + ) -> ndarray | Any: ... + @overload def multivariate_normal[MeanShape: IntTuple, CovShape: IntTuple]( self, /, diff --git a/tensor-shapes/pyrefly-numpy-stubs/test/test_random.py b/tensor-shapes/pyrefly-numpy-stubs/test/test_random.py index 64d07b59dc..30a4acbd98 100644 --- a/tensor-shapes/pyrefly-numpy-stubs/test/test_random.py +++ b/tensor-shapes/pyrefly-numpy-stubs/test/test_random.py @@ -269,3 +269,83 @@ def check_random_and_standard_cauchy_symbolic_size[N: IntVar, M: IntVar]( rng = np.random.default_rng(0) assert_type(rng.random(n), np.ndarray[[N], np.dtype[np.float64]]) assert_type(rng.standard_cauchy((n, m)), np.ndarray[[N, M], np.dtype[np.float64]]) + + +def check_integers_return_scalars_without_size() -> None: + rng = np.random.default_rng(0) + assert_type(rng.integers(10), np.int64) + assert_type(rng.integers(0, 10), np.int64) + assert_type(rng.integers(0, 10, dtype=np.uint8), np.uint8) + + +def test_integers_int_size() -> None: + rng = np.random.default_rng(0) + result = rng.integers(0, 10, size=5) + assert_type(result, np.ndarray[[5], np.dtype[np.int64]]) + assert_shape(result.shape, (5,)) + assert_shape(rng.integers(0, 10, 5).shape, (5,)) + assert_shape(rng.integers(10, size=5).shape, (5,)) + assert_shape(rng.integers(0, 10, 5, endpoint=True).shape, (5,)) + + +def test_integers_tuple_size() -> None: + rng = np.random.default_rng(0) + result = rng.integers(0, 10, (2, 3)) + assert_type(result, np.ndarray[[2, 3], np.dtype[np.int64]]) + assert_shape(result.shape, (2, 3)) + assert_shape(rng.integers(10, size=(2, 3)).shape, (2, 3)) + + +def test_integers_dtype() -> None: + rng = np.random.default_rng(0) + result = rng.integers(0, 10, (2, 3), dtype=np.int32) + assert_type(result, np.ndarray[[2, 3], np.dtype[np.int32]]) + assert_shape(result.shape, (2, 3)) + mask = rng.integers(0, 2, size=4, dtype=np.bool_) + assert_type(mask, np.ndarray[[4], np.dtype[np.bool_]]) + assert_shape(mask.shape, (4,)) + named = rng.integers(0, 10, size=4, dtype="int64") + assert_type(named, np.ndarray[[4], np.dtype[np.int64]]) + assert_shape(named.shape, (4,)) + + +def test_integers_array_bounds_with_size() -> None: + rng = np.random.default_rng(0) + result = rng.integers(np.full(3, 5), size=(2, 3)) + assert_type(result, np.ndarray[[2, 3], np.dtype[np.int64]]) + assert_shape(result.shape, (2, 3)) + + +def test_integers_array_bounds_without_size() -> None: + rng = np.random.default_rng(0) + result = rng.integers(np.zeros(3, dtype=np.int64), np.full((2, 1), 5)) + assert_type(result, np.ndarray[[2, 3], np.dtype[np.int64]]) + assert_shape(result.shape, (2, 3)) + assert_shape(rng.integers(np.full(3, 5)).shape, (3,)) + assert_shape(rng.integers(0, np.full((2, 1), 5)).shape, (2, 1)) + + +def test_integers_unshaped_arguments() -> None: + rng = np.random.default_rng(0) + assert_shape(rng.integers(10, size=[2, 3]).shape, IntTuple, runtime=(2, 3)) + assert_shape( + rng.integers(10, size=(2, 3), dtype=int).shape, IntTuple, runtime=(2, 3) + ) + + +def test_integers_mismatched_bounds() -> None: + rng = np.random.default_rng(0) + with assert_raises(ValueError): + rng.integers(np.full(3, 5), np.full(2, 9)) # E: Cannot broadcast + + +def check_integers_symbolic_shapes[N: IntVar, M: IntVar]( + n: Int[N], low: np.ndarray[[N, 1]], high: np.ndarray[[M]] +) -> None: + rng = np.random.default_rng(0) + assert_type(rng.integers(10, size=n), np.ndarray[[N], np.dtype[np.int64]]) + assert_type(rng.integers(low, high), np.ndarray[[N, M], np.dtype[np.int64]]) + assert_type( + rng.integers(10, size=(n, n), dtype=np.int16), + np.ndarray[[N, N], np.dtype[np.int16]], + )