diff --git a/tensor-shapes/pyrefly-numpy-stubs/numpy-stubs/__init__.pyi b/tensor-shapes/pyrefly-numpy-stubs/numpy-stubs/__init__.pyi index 7a32fce55e..4a1e31ed23 100644 --- a/tensor-shapes/pyrefly-numpy-stubs/numpy-stubs/__init__.pyi +++ b/tensor-shapes/pyrefly-numpy-stubs/numpy-stubs/__init__.pyi @@ -6,7 +6,17 @@ from builtins import bool as py_bool from collections.abc import Iterator, Sequence from types import EllipsisType -from typing import Any, Final, Generic, Literal, Never, overload, SupportsIndex, TypeVar +from typing import ( + Any, + Final, + Generic, + Literal, + Never, + overload, + Self, + SupportsIndex, + TypeVar, +) import shape_extensions from numpy.__config__ import ( @@ -206,7 +216,7 @@ from numpy._shapes import ( vecdot_shape, vecmat_shape, ) -from numpy._typing import _ArrayLikeInt_co, ArrayLike +from numpy._typing import _ArrayLikeInt_co, ArrayLike, DTypeLike from numpy._typing._extended_precision import ( complex256 as complex256, float128 as float128, @@ -429,25 +439,75 @@ _ShapeT_co = TypeVar("_ShapeT_co", bound=_Shape, default=_Shape, covariant=True) type _Axis = int | tuple[int, ...] | None type _SingleAxis = int | None type _BasicIndex = int | slice | list[int] | None | EllipsisType +type _NumberScalar = number | bool_ | int | float | complex class generic: + @property + def shape(self) -> tuple[()]: ... + def __bool__(self, /) -> py_bool: ... + @overload + def __array__(self, dtype: None = None, /) -> ndarray[[], _dtype[Self]]: ... + @overload + def __array__[DType: _dtype](self, dtype: DType, /) -> ndarray[[], DType]: ... def astype(self, dtype: Any, **kwargs: Any) -> Any: ... -class number(generic): ... -class integer(number): ... +class number(generic): + def __neg__(self) -> Self: ... + def __pos__(self) -> Self: ... + def __abs__(self) -> number: ... + # The result dtype depends on both operands. + def __add__(self, other: _NumberScalar, /) -> number: ... + def __radd__(self, other: _NumberScalar, /) -> number | complex: ... + def __sub__(self, other: _NumberScalar, /) -> number: ... + def __rsub__(self, other: _NumberScalar, /) -> number | complex: ... + def __mul__(self, other: _NumberScalar, /) -> number: ... + def __rmul__(self, other: _NumberScalar, /) -> number | complex: ... + def __truediv__(self, other: _NumberScalar, /) -> number: ... + def __rtruediv__(self, other: _NumberScalar, /) -> number | complex: ... + def __pow__(self, other: _NumberScalar, mod: None = None, /) -> number: ... + def __rpow__( + self, other: _NumberScalar, mod: None = None, / + ) -> number | complex: ... + def __lt__(self, other: _NumberScalar, /) -> bool_: ... + def __le__(self, other: _NumberScalar, /) -> bool_: ... + def __gt__(self, other: _NumberScalar, /) -> bool_: ... + def __ge__(self, other: _NumberScalar, /) -> bool_: ... + def __int__(self, /) -> int: ... + def __float__(self, /) -> float: ... -class signedinteger(integer): +class integer(number): + def __index__(self, /) -> int: ... # TODO(stroxler): Sequence inputs to scalar constructors can produce arrays. def __init__(self, value: object = ...) -> None: ... +class signedinteger(integer): ... class unsignedinteger(integer): ... class inexact(number): ... class floating(inexact): + def __init__( + self, value: number | bool_ | int | float | str | bytes = ... + ) -> None: ... + +class complexfloating(inexact): + def __init__(self, value: _NumberScalar | str | bytes = ...) -> None: ... + def __complex__(self, /) -> complex: ... + +class bool_(generic): + def __init__(self, value: _NumberScalar | str | bytes = ...) -> None: ... + def __int__(self, /) -> int: ... def __float__(self, /) -> float: ... + def __add__(self, other: _NumberScalar, /) -> number | bool_: ... + def __radd__(self, other: _NumberScalar, /) -> number | bool_: ... + def __sub__(self, other: number | int | float | complex, /) -> number: ... + def __rsub__(self, other: number | int | float | complex, /) -> number: ... + def __mul__(self, other: _NumberScalar, /) -> number | bool_: ... + def __rmul__(self, other: _NumberScalar, /) -> number | bool_: ... + def __truediv__(self, other: _NumberScalar, /) -> number: ... + def __rtruediv__(self, other: _NumberScalar, /) -> number: ... + def __pow__(self, other: _NumberScalar, mod: None = None, /) -> number: ... + def __rpow__(self, other: _NumberScalar, mod: None = None, /) -> number: ... -class complexfloating(inexact): ... -class bool_(generic): ... class float16(floating): ... class float32(floating): ... class float64(floating): ... @@ -592,6 +652,7 @@ class ndarray(Generic[_ShapeT_co, _DTypeT_co]): @overload def __iter__[N: IntVar](self: ndarray[[N], _DTypeT_co], /) -> Iterator[Any]: ... def __int__(self: ndarray[[]], /) -> int: ... + def __index__(self: ndarray[[], _dtype[integer]], /) -> int: ... def __float__(self: ndarray[[]], /) -> float: ... def __complex__(self: ndarray[[]], /) -> complex: ... # Only 2-D transpose is modeled for the NumPy shape-stub MVP. @@ -600,30 +661,54 @@ class ndarray(Generic[_ShapeT_co, _DTypeT_co]): self: ndarray[[N, P], _DTypeT_co], ) -> ndarray[[P, N], _DTypeT_co]: ... # Unpacking preserves gradual array shapes instead of using the scalar default. + @overload def __add__[OtherShape: _Shape = []]( self, other: ndarray[[*OtherShape]] | int | float ) -> ndarray[_broadcast_shape(_ShapeT_co, OtherShape), _DTypeT_co]: ... + @overload + def __add__(self, other: complex | number | bool_) -> ndarray[_ShapeT_co]: ... + @overload def __radd__[OtherShape: _Shape = []]( self, other: ndarray[[*OtherShape]] | int | float ) -> ndarray[_broadcast_shape(_ShapeT_co, OtherShape), _DTypeT_co]: ... + @overload + def __radd__(self, other: complex | number | bool_) -> ndarray[_ShapeT_co]: ... + @overload def __sub__[OtherShape: _Shape = []]( self, other: ndarray[[*OtherShape]] | int | float ) -> ndarray[_broadcast_shape(_ShapeT_co, OtherShape), _DTypeT_co]: ... + @overload + def __sub__(self, other: complex | number | bool_) -> ndarray[_ShapeT_co]: ... + @overload def __rsub__[OtherShape: _Shape = []]( self, other: ndarray[[*OtherShape]] | int | float ) -> ndarray[_broadcast_shape(_ShapeT_co, OtherShape), _DTypeT_co]: ... + @overload + def __rsub__(self, other: complex | number | bool_) -> ndarray[_ShapeT_co]: ... + @overload def __truediv__[OtherShape: _Shape = []]( self, other: ndarray[[*OtherShape]] | int | float ) -> ndarray[_broadcast_shape(_ShapeT_co, OtherShape), _DTypeT_co]: ... + @overload + def __truediv__(self, other: complex | number | bool_) -> ndarray[_ShapeT_co]: ... + @overload def __rtruediv__[OtherShape: _Shape = []]( self, other: ndarray[[*OtherShape]] | int | float ) -> ndarray[_broadcast_shape(_ShapeT_co, OtherShape), _DTypeT_co]: ... + @overload + def __rtruediv__(self, other: complex | number | bool_) -> ndarray[_ShapeT_co]: ... + @overload def __mul__[OtherShape: _Shape = []]( self, other: ndarray[[*OtherShape]] | int | float ) -> ndarray[_broadcast_shape(_ShapeT_co, OtherShape), _DTypeT_co]: ... + @overload + def __mul__(self, other: complex | number | bool_) -> ndarray[_ShapeT_co]: ... + @overload def __rmul__[OtherShape: _Shape = []]( self, other: ndarray[[*OtherShape]] | int | float ) -> ndarray[_broadcast_shape(_ShapeT_co, OtherShape), _DTypeT_co]: ... + @overload + def __rmul__(self, other: complex | number | bool_) -> ndarray[_ShapeT_co]: ... def __floordiv__[OtherShape: _Shape = []]( self, other: ndarray[[*OtherShape]] | int | float ) -> ndarray[_broadcast_shape(_ShapeT_co, OtherShape), _DTypeT_co]: ... @@ -639,12 +724,18 @@ class ndarray(Generic[_ShapeT_co, _DTypeT_co]): def __neg__(self) -> ndarray[_ShapeT_co, _DTypeT_co]: ... def __pos__(self) -> ndarray[_ShapeT_co, _DTypeT_co]: ... def __abs__(self, /) -> ndarray[_ShapeT_co]: ... + @overload def __pow__[OtherShape: _Shape = []]( self, other: ndarray[[*OtherShape]] | int | float ) -> ndarray[_broadcast_shape(_ShapeT_co, OtherShape), _DTypeT_co]: ... + @overload + def __pow__(self, other: complex | number | bool_) -> ndarray[_ShapeT_co]: ... + @overload def __rpow__[OtherShape: _Shape = []]( self, other: ndarray[[*OtherShape]] | int | float ) -> ndarray[_broadcast_shape(_ShapeT_co, OtherShape), _DTypeT_co]: ... + @overload + def __rpow__(self, other: complex | number | bool_) -> ndarray[_ShapeT_co]: ... # Scalar results also satisfy Python's truth-valued comparison protocols. @overload def __lt__(self: ndarray[[]], other: Any, /) -> Any: ... @@ -745,12 +836,14 @@ class ndarray(Generic[_ShapeT_co, _DTypeT_co]): axis: Axis = None, *, keepdims: KeepDims = False, + initial: _ArrayScalar = ..., ) -> ndarray[reduce_shape(_ShapeT_co, Axis, KeepDims), _DTypeT_co]: ... def max[Axis: Flag[_Axis], KeepDims: Flag[py_bool]]( self, axis: Axis = None, *, keepdims: KeepDims = False, + initial: _ArrayScalar = ..., ) -> ndarray[reduce_shape(_ShapeT_co, Axis, KeepDims), _DTypeT_co]: ... def all[Axis: Flag[_Axis], KeepDims: Flag[py_bool]]( self, @@ -1626,7 +1719,7 @@ def asarray( copy: py_bool | None = ..., like: Any, ) -> Any: ... -def round[Shape: _Shape](x: ndarray[Shape]) -> ndarray[Shape]: ... +def round[Shape: _Shape](a: ndarray[Shape], decimals: int = 0) -> ndarray[Shape]: ... # Interpolation tracks query shapes and leaves the result dtype gradual. @overload @@ -1676,7 +1769,25 @@ def diag[M: IntVar, N: IntVar, DType, K: Flag[int] = 0]( def diag[S: _Shape, DType]( v: ndarray[S, DType], k: int = 0 ) -> ndarray[IntTuple, DType]: ... -def arange[N: IntVar](stop: Int[N], /) -> ndarray[[N], dtype[intp]]: ... +@overload +def arange[N: IntVar]( + stop: Int[N], + /, + *, + dtype: None = None, + device: Literal["cpu"] | None = None, + like: None = None, +) -> ndarray[[N], dtype[intp]]: ... +@overload +def arange( + start: float, + stop: float | None = None, + step: float = 1, + dtype: DTypeLike | None = None, + *, + device: Literal["cpu"] | None = None, + like: None = None, +) -> ndarray[[int]]: ... def expand_dims[Shape: _Shape, DType, Axis: Flag[int]]( a: ndarray[Shape, DType], axis: Axis ) -> ndarray[expand_dims_shape(Shape, Axis), DType]: ... @@ -1705,11 +1816,19 @@ def min[ Axis: Flag[_Axis], KeepDims: Flag[py_bool], ]( - a: ndarray[Shape, DType], axis: Axis = None, *, keepdims: KeepDims = False + a: ndarray[Shape, DType], + axis: Axis = None, + *, + keepdims: KeepDims = False, + initial: _ArrayScalar = ..., ) -> ndarray[reduce_shape(Shape, Axis, KeepDims), Any]: ... @overload def min( - a: Sequence[ArrayLike], axis: _Axis = None, *, keepdims: py_bool = False + a: Sequence[ArrayLike], + axis: _Axis = None, + *, + keepdims: py_bool = False, + initial: _ArrayScalar = ..., ) -> Any: ... @overload def max[ @@ -1718,11 +1837,19 @@ def max[ Axis: Flag[_Axis], KeepDims: Flag[py_bool], ]( - a: ndarray[Shape, DType], axis: Axis = None, *, keepdims: KeepDims = False + a: ndarray[Shape, DType], + axis: Axis = None, + *, + keepdims: KeepDims = False, + initial: _ArrayScalar = ..., ) -> ndarray[reduce_shape(Shape, Axis, KeepDims), Any]: ... @overload def max( - a: Sequence[ArrayLike], axis: _Axis = None, *, keepdims: py_bool = False + a: Sequence[ArrayLike], + axis: _Axis = None, + *, + keepdims: py_bool = False, + initial: _ArrayScalar = ..., ) -> Any: ... @overload def argmin[ diff --git a/tensor-shapes/pyrefly-numpy-stubs/test/test_arithmetic.py b/tensor-shapes/pyrefly-numpy-stubs/test/test_arithmetic.py index ea2deee466..43a3d93d5c 100644 --- a/tensor-shapes/pyrefly-numpy-stubs/test/test_arithmetic.py +++ b/tensor-shapes/pyrefly-numpy-stubs/test/test_arithmetic.py @@ -53,6 +53,31 @@ def test_unary_arithmetic() -> None: assert_shape((-np.ones((3, 4))).shape, (3, 4)) +def test_scalar_array_arithmetic() -> None: + a = np.full((2, 3), 2.0, dtype=np.float32) + + assert_shape((a + 1j).shape, (2, 3)) + assert_shape((1j + a).shape, (2, 3)) + assert_shape((a - 1j).shape, (2, 3)) + assert_shape((1j - a).shape, (2, 3)) + assert_shape((a * 1j).shape, (2, 3)) + assert_shape((1j * a).shape, (2, 3)) + assert_shape((a / 1j).shape, (2, 3)) + assert_shape((1j / a).shape, (2, 3)) + assert_shape((a**1j).shape, (2, 3)) + assert_shape((1j**a).shape, (2, 3)) + assert_type(a + 1j, np.ndarray[[2, 3]]) + assert (a + 1j).dtype == np.complex64 + + scalar = np.array(2.0) + assert_shape((scalar * 1j).shape, ()) + + real = np.float64(2.0) + assert_shape((real * a).shape, (2, 3)) + assert_shape((a * real).shape, (2, 3)) + assert (a * real).dtype == np.float64 + + def test_builtin_abs_preserves_shape() -> None: scalar = np.array(-2.0) vector = np.full(5, -1.0) diff --git a/tensor-shapes/pyrefly-numpy-stubs/test/test_dtype_properties.py b/tensor-shapes/pyrefly-numpy-stubs/test/test_dtype_properties.py index d80f8c0b17..8a6ad3710e 100644 --- a/tensor-shapes/pyrefly-numpy-stubs/test/test_dtype_properties.py +++ b/tensor-shapes/pyrefly-numpy-stubs/test/test_dtype_properties.py @@ -8,7 +8,7 @@ from typing import Any, assert_type import numpy as np -from shape_extensions import assert_shape +from shape_extensions import assert_raises, assert_shape def check_number_is_not_gradual() -> None: @@ -64,6 +64,62 @@ def test_floating_and_complex_scalar_hierarchy() -> None: assert np.cdouble is np.complex128 +def test_numeric_scalar_operations() -> None: + integer = np.int64(2) + real = np.float64(2.0) + complex_ = np.complex128(1.0 + 2.0j) + + assert_type(int(integer), int) + assert_type(integer.__index__(), int) + assert_type(float(real), float) + assert_type(complex(complex_), complex) + assert_shape((real**2).shape, ()) + assert_shape((real * 1j).shape, ()) + assert_type(1j * real, np.number | complex) + assert_shape((complex_ / integer).shape, ()) + assert_shape(abs(complex_).shape, ()) + assert int(np.uint64(2)) == 2 + assert [10, 20, 30][integer] == 30 + assert_type(max([real, real]), np.float64) + assert max([real, real + 1]) == 3.0 + scalar_array = real.__array__() + assert_shape(scalar_array.shape, ()) + assert_type(scalar_array.dtype, np.dtype[np.float64]) + converted = real.__array__(np.dtype(np.complex128)) + assert_shape(converted.shape, ()) + assert_type(converted.dtype, np.dtype[np.complex128]) + + boolean = np.bool_(True) + assert_type(int(boolean), int) + assert_shape((boolean * 2).shape, ()) + assert_shape((boolean**2).shape, ()) + assert_shape((1j * boolean).shape, ()) + assert boolean - 2.0 == -1.0 + assert 2.0 - boolean == 1.0 + + with assert_raises(TypeError): + real + "invalid" # E: not supported + + with assert_raises(TypeError): + boolean - boolean # E: not supported + + with assert_raises(TypeError): + np.float64(1j) # E: not assignable + + +def test_integer_scalar_array_index() -> None: + index = np.ones((), dtype=np.int64) + assert_shape(index.shape, ()) + assert_type(index.__index__(), int) + assert [10, 20][index] == 20 + + with assert_raises(TypeError): + np.ones((), dtype=np.float64).__index__() # E: not assignable + + with assert_raises(TypeError): + np.ones((1,), dtype=np.int64).__index__() # E: not assignable + + def test_flexible_and_datetime_scalar_hierarchy() -> None: assert_shape(np.zeros((2,), dtype=np.object_).shape, (2,)) assert_type(np.flexible, type[np.flexible]) diff --git a/tensor-shapes/pyrefly-numpy-stubs/test/test_indexing.py b/tensor-shapes/pyrefly-numpy-stubs/test/test_indexing.py index 6acb7db7e0..cf45fb9c89 100644 --- a/tensor-shapes/pyrefly-numpy-stubs/test/test_indexing.py +++ b/tensor-shapes/pyrefly-numpy-stubs/test/test_indexing.py @@ -29,6 +29,16 @@ def test_arange_from_array_length() -> None: assert indices.dtype == np.dtype(np.intp) +def test_arange_optional_arguments() -> None: + assert_shape(np.arange(2, 8).shape, (int,), runtime=(6,)) + assert_shape(np.arange(2, 9, 2).shape, (int,), runtime=(4,)) + assert_shape(np.arange(0.5, 2.0, 0.5).shape, (int,), runtime=(3,)) + assert_shape(np.arange(2.5).shape, (int,), runtime=(3,)) + values = np.arange(4, dtype=np.float32, device="cpu") + assert_shape(values.shape, (int,), runtime=(4,)) + assert values.dtype == np.float32 + + def test_array_protocol_shapes() -> None: vector = np.ones((5,)) matrix = np.ones((5, 3)) diff --git a/tensor-shapes/pyrefly-numpy-stubs/test/test_math_ufuncs.py b/tensor-shapes/pyrefly-numpy-stubs/test/test_math_ufuncs.py index fdff53c707..7a2fd87c3b 100644 --- a/tensor-shapes/pyrefly-numpy-stubs/test/test_math_ufuncs.py +++ b/tensor-shapes/pyrefly-numpy-stubs/test/test_math_ufuncs.py @@ -234,5 +234,7 @@ def test_rounding_preserves_shape() -> None: assert_shape(np.floor(a).shape, (5,)) assert_shape(np.ceil(a).shape, (5,)) assert_shape(np.round(a).shape, (5,)) + assert_shape(np.round(a, 2).shape, (5,)) + assert_shape(np.round(a=a, decimals=-1).shape, (5,)) assert_shape(np.trunc(a).shape, (5,)) assert_shape(np.clip(a, -1.0, 2.0).shape, (5,)) diff --git a/tensor-shapes/pyrefly-numpy-stubs/test/test_reductions.py b/tensor-shapes/pyrefly-numpy-stubs/test/test_reductions.py index ca91836484..f3bc8b667a 100644 --- a/tensor-shapes/pyrefly-numpy-stubs/test/test_reductions.py +++ b/tensor-shapes/pyrefly-numpy-stubs/test/test_reductions.py @@ -22,6 +22,16 @@ def test_reduce_matrix_no_axis() -> None: assert_shape(np.mean(a).shape, ()) +def test_extrema_initial_value() -> None: + a = np.zeros((0, 3)) + + assert_shape(a.max(axis=0, initial=0.0).shape, (3,)) + assert_shape(a.min(axis=0, initial=0.0).shape, (3,)) + assert_shape(np.max(a, axis=0, initial=0.0).shape, (3,)) + assert_shape(np.min(a, axis=0, initial=0.0).shape, (3,)) + assert np.array_equal(np.max(a, axis=0, initial=2.0), np.full(3, 2.0)) + + def test_reduce_higher_rank_no_axis() -> None: a = cast( "np.ndarray[tuple[Literal[2], Literal[3], Literal[4]]]",