Infer NumPy Generator.normal result shapes

tensor-shapes/pyrefly-numpy-stubs/numpy-stubs/random/__init__.pyi

@@ -3,20 +3,27 @@
 # 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,
     _DTypeLikeF32,
     _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

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