From 05463c2ab8aa79e3e65600bce025455836154a01 Mon Sep 17 00:00:00 2001 From: AlessandroRuggiero Date: Sat, 14 Nov 2026 11:46:36 +0100 Subject: [PATCH] Infer NumPy Generator.uniform and exponential result shapes --- .../numpy-stubs/random/__init__.pyi | 126 ++++++++++++++++++ .../test/test_issue_4917.py | 15 +-- .../pyrefly-numpy-stubs/test/test_random.py | 78 +++++++++++ 3 files changed, 210 insertions(+), 9 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 c3775cd5f3..9a3c5297a7 100644 --- a/tensor-shapes/pyrefly-numpy-stubs/numpy-stubs/random/__init__.pyi +++ b/tensor-shapes/pyrefly-numpy-stubs/numpy-stubs/random/__init__.pyi @@ -426,6 +426,132 @@ class Generator(_NumpyGenerator): scale: _ArrayLikeFloat_co = 1.0, size: None = None, ) -> _ArrayF64 | Any: ... + @overload + def uniform( + self, /, low: _FloatLike_co = 0.0, high: _FloatLike_co = 1.0, size: None = None + ) -> float: ... + # With `size`, the result has exactly that shape, and NumPy raises if `low` or + # `high` does not broadcast to it. + @overload + def uniform[N: IntVar]( + self, /, low: _ArrayLikeFloat_co, high: _ArrayLikeFloat_co, size: Int[N] + ) -> ndarray[[N], dtype[float64]]: ... + @overload + def uniform[Shape: IntTuple]( + self, /, low: _ArrayLikeFloat_co, high: _ArrayLikeFloat_co, size: Shape + ) -> ndarray[Shape, dtype[float64]]: ... + @overload + def uniform[N: IntVar]( + self, + /, + low: _ArrayLikeFloat_co = 0.0, + high: _ArrayLikeFloat_co = 1.0, + *, + size: Int[N], + ) -> ndarray[[N], dtype[float64]]: ... + @overload + def uniform[Shape: IntTuple]( + self, + /, + low: _ArrayLikeFloat_co = 0.0, + high: _ArrayLikeFloat_co = 1.0, + *, + size: Shape, + ) -> ndarray[Shape, dtype[float64]]: ... + # Without `size`, the result has the broadcast shape of `low` and `high`. + @overload + def uniform[LowShape: IntTuple = [], HighShape: IntTuple = []]( + self, + /, + low: ndarray[LowShape] | float = 0.0, + high: ndarray[HighShape] | float = 1.0, + size: None = None, + ) -> ndarray[broadcast(LowShape, HighShape), dtype[float64]]: ... + @overload + def uniform( + self, /, low: _ArrayLikeFloat_co, high: _ArrayLikeFloat_co, size: _ShapeLike + ) -> _ArrayF64: ... + @overload + def uniform( + self, /, low: _ArrayLikeFloat_co, high: _NDArrayLikeFloat, size: None = None + ) -> _ArrayF64: ... + @overload + def uniform( + self, + /, + low: _ArrayLikeFloat_co = 0.0, + high: _ArrayLikeFloat_co = 1.0, + *, + size: _ShapeLike, + ) -> _ArrayF64: ... + @overload + def uniform( + self, + /, + low: _ArrayLikeFloat_co = 0.0, + *, + high: _NDArrayLikeFloat, + size: None = None, + ) -> _ArrayF64: ... + @overload + def uniform( + self, + /, + low: _NDArrayLikeFloat, + high: _ArrayLikeFloat_co = 1.0, + size: None = None, + ) -> _ArrayF64: ... + @overload + def uniform( + self, + /, + low: _ArrayLikeFloat_co = 0.0, + high: _ArrayLikeFloat_co = 1.0, + size: None = None, + ) -> _ArrayF64 | Any: ... + @overload + def exponential( + self, /, scale: _FloatLike_co = 1.0, size: None = None + ) -> float: ... + # With `size`, the result has exactly that shape, and NumPy raises if `scale` + # does not broadcast to it. + @overload + def exponential[N: IntVar]( + self, /, scale: _ArrayLikeFloat_co, size: Int[N] + ) -> ndarray[[N], dtype[float64]]: ... + @overload + def exponential[Shape: IntTuple]( + self, /, scale: _ArrayLikeFloat_co, size: Shape + ) -> ndarray[Shape, dtype[float64]]: ... + @overload + def exponential[N: IntVar]( + self, /, scale: _ArrayLikeFloat_co = 1.0, *, size: Int[N] + ) -> ndarray[[N], dtype[float64]]: ... + @overload + def exponential[Shape: IntTuple]( + self, /, scale: _ArrayLikeFloat_co = 1.0, *, size: Shape + ) -> ndarray[Shape, dtype[float64]]: ... + # Without `size`, the result has the shape of `scale`. + @overload + def exponential[Shape: IntTuple]( + self, /, scale: ndarray[Shape], size: None = None + ) -> ndarray[Shape, dtype[float64]]: ... + @overload + def exponential( + self, /, scale: _ArrayLikeFloat_co = 1.0, *, size: _ShapeLike + ) -> _ArrayF64: ... + @overload + def exponential( + self, /, scale: _ArrayLikeFloat_co, size: _ShapeLike + ) -> _ArrayF64: ... + @overload + def exponential( + self, /, scale: _NDArrayLikeFloat, size: None = None + ) -> _ArrayF64: ... + @overload + def exponential( + self, /, scale: _ArrayLikeFloat_co, size: None = None + ) -> _ArrayF64 | Any: ... def default_rng( seed: _ArrayLikeInt_co diff --git a/tensor-shapes/pyrefly-numpy-stubs/test/test_issue_4917.py b/tensor-shapes/pyrefly-numpy-stubs/test/test_issue_4917.py index 7cee05f7cd..9f375ab035 100644 --- a/tensor-shapes/pyrefly-numpy-stubs/test/test_issue_4917.py +++ b/tensor-shapes/pyrefly-numpy-stubs/test/test_issue_4917.py @@ -5,10 +5,10 @@ from __future__ import annotations -from typing import Any, assert_type +from typing import assert_type import numpy as np -from shape_extensions import assert_shape, IntTuple +from shape_extensions import assert_shape def test_pairwise_distance_comparison() -> None: @@ -22,14 +22,11 @@ def test_pairwise_distance_comparison() -> None: deltas = positions[:, None, :] - positions[None, :, :] deltas -= size * np.round(deltas / size) squared_distances = (deltas**2).sum(-1) - # `deltas` has a gradual rank, so `np.round` may produce a scalar and returns `Any`. - assert_type(squared_distances, Any) + assert_type(squared_distances, np.ndarray[[300, 300], np.dtype[np.float64]]) assert isinstance(squared_distances, np.ndarray) - assert_shape(squared_distances.shape, IntTuple, runtime=(count, count)) + assert_shape(squared_distances.shape, (300, 300)) assert squared_distances.dtype == np.dtype(np.float64) nearby = squared_distances < radius * radius - - # Comparisons with unknown shapes also match the scalar overload returning Any. - assert_type(nearby, Any) - assert_shape(nearby.shape, IntTuple, runtime=(count, count)) + assert_type(nearby, np.ndarray[[300, 300], np.dtype[np.bool_]]) + assert_shape(nearby.shape, (300, 300)) assert nearby.dtype == np.bool_ diff --git a/tensor-shapes/pyrefly-numpy-stubs/test/test_random.py b/tensor-shapes/pyrefly-numpy-stubs/test/test_random.py index 64d07b59dc..95d34b726b 100644 --- a/tensor-shapes/pyrefly-numpy-stubs/test/test_random.py +++ b/tensor-shapes/pyrefly-numpy-stubs/test/test_random.py @@ -153,6 +153,84 @@ def check_normal_symbolic_shapes[N: IntVar, M: IntVar]( assert_type(rng.normal(loc, scale), np.ndarray[[N, M], np.dtype[np.float64]]) +def check_uniform_and_exponential_return_float_without_size() -> None: + rng = np.random.default_rng(0) + assert_type(rng.uniform(), float) + assert_type(rng.uniform(1.0, 2.0), float) + assert_type(rng.exponential(), float) + assert_type(rng.exponential(2.0), float) + + +def test_uniform_and_exponential_int_size() -> None: + rng = np.random.default_rng(0) + result = rng.uniform(size=5) + assert_type(result, np.ndarray[[5], np.dtype[np.float64]]) + assert_shape(result.shape, (5,)) + assert_shape(rng.uniform(0.0, 1.0, 5).shape, (5,)) + assert_shape(rng.exponential(size=5).shape, (5,)) + assert_shape(rng.exponential(2.0, 5).shape, (5,)) + + +def test_uniform_and_exponential_tuple_size() -> None: + rng = np.random.default_rng(0) + result = rng.exponential(2.0, (2, 3)) + assert_type(result, np.ndarray[[2, 3], np.dtype[np.float64]]) + assert_shape(result.shape, (2, 3)) + assert_shape(rng.exponential(size=(2, 3)).shape, (2, 3)) + assert_shape(rng.uniform(0.0, 1.0, (2, 3)).shape, (2, 3)) + assert_shape(rng.uniform(size=(2, 3)).shape, (2, 3)) + + +def test_uniform_and_exponential_array_parameters_with_size() -> None: + rng = np.random.default_rng(0) + result = rng.uniform(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.exponential(np.ones(3), size=(2, 3)).shape, (2, 3)) + + +def test_uniform_array_parameters_without_size() -> None: + rng = np.random.default_rng(0) + result = rng.uniform(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.uniform(np.zeros(3)).shape, (3,)) + assert_shape(rng.uniform(high=np.ones((2, 1))).shape, (2, 1)) + + +def test_exponential_array_scale_without_size() -> None: + rng = np.random.default_rng(0) + result = rng.exponential(np.ones((2, 3))) + assert_type(result, np.ndarray[[2, 3], np.dtype[np.float64]]) + assert_shape(result.shape, (2, 3)) + + +def test_uniform_and_exponential_list_size_and_parameters() -> None: + rng = np.random.default_rng(0) + assert_shape(rng.uniform(size=[2, 3]).shape, IntTuple, runtime=(2, 3)) + assert_shape(rng.exponential(size=[2, 3]).shape, IntTuple, runtime=(2, 3)) + assert_shape(rng.exponential([1.0, 2.0]).shape, IntTuple, runtime=(2,)) + + +def test_uniform_mismatched_parameters() -> None: + rng = np.random.default_rng(0) + with assert_raises(ValueError): + rng.uniform(np.zeros(3), np.ones(2)) # E: Cannot broadcast + + +def check_uniform_and_exponential_symbolic_shapes[N: IntVar, M: IntVar]( + n: Int[N], + low: np.ndarray[[N, 1]], + high: np.ndarray[[M]], + scale: np.ndarray[[N, M]], +) -> None: + rng = np.random.default_rng(0) + assert_type(rng.uniform(size=n), np.ndarray[[N], np.dtype[np.float64]]) + assert_type(rng.uniform(low, high), np.ndarray[[N, M], np.dtype[np.float64]]) + assert_type(rng.exponential(size=n), np.ndarray[[N], np.dtype[np.float64]]) + assert_type(rng.exponential(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))