Infer NumPy Generator.uniform and exponential result shapes

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

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_

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))