diff --git a/tensor-shapes/jax-pyrefly-stubs/jax-stubs/_array.pyi b/tensor-shapes/jax-pyrefly-stubs/jax-stubs/_array.pyi index 954aa73cb1..fc72587524 100644 --- a/tensor-shapes/jax-pyrefly-stubs/jax-stubs/_array.pyi +++ b/tensor-shapes/jax-pyrefly-stubs/jax-stubs/_array.pyi @@ -131,10 +131,10 @@ class Array[Shape: _Shape = _Shape]: def __len__[N: IntVar, Rest: _Shape = []]( self: Array[[N, *Rest]], ) -> Int[N]: ... - def __iter__[N: IntVar, Rest: _Shape = []]( + def __iter__[N: IntVar, Rest: _Shape]( self: Array[[N, *Rest]], / ) -> Iterator[Array[Rest]]: ... - def __reversed__[N: IntVar, Rest: _Shape = []]( + def __reversed__[N: IntVar, Rest: _Shape]( self: Array[[N, *Rest]], / ) -> Iterator[Array[Rest]]: ... def __contains__(self, key: object) -> bool: ... diff --git a/tensor-shapes/jax-pyrefly-stubs/test/test_typing.py b/tensor-shapes/jax-pyrefly-stubs/test/test_typing.py index 79f6dbb39d..9bf058a33f 100644 --- a/tensor-shapes/jax-pyrefly-stubs/test/test_typing.py +++ b/tensor-shapes/jax-pyrefly-stubs/test/test_typing.py @@ -19,6 +19,17 @@ def asarray(x: ArrayLike) -> jax.Array: return jnp.asarray(x) +def test_iteration_preserves_unknown_rank() -> None: + def rows(x: jax.Array) -> None: + assert_shape(next(iter(x)).shape, IntTuple, runtime=(3,)) + assert_shape(next(reversed(x)).shape, IntTuple, runtime=(3,)) + + rows(jnp.ones((2, 3))) + assert_shape(next(iter(jnp.ones(3))).shape, ()) + assert_shape(next(iter(jnp.ones((2, 3)))).shape, (3,)) + assert_shape(next(reversed(jnp.ones((2, 3)))).shape, (3,)) + + def rejects_sequence(x: list[int]) -> None: asarray(x) # E: not assignable to parameter `x` diff --git a/tensor-shapes/pyrefly-numpy-stubs/numpy-stubs/__init__.pyi b/tensor-shapes/pyrefly-numpy-stubs/numpy-stubs/__init__.pyi index 4a1e31ed23..5bbc3a4575 100644 --- a/tensor-shapes/pyrefly-numpy-stubs/numpy-stubs/__init__.pyi +++ b/tensor-shapes/pyrefly-numpy-stubs/numpy-stubs/__init__.pyi @@ -644,13 +644,13 @@ class ndarray(Generic[_ShapeT_co, _DTypeT_co]): ) -> ndarray[IntTuple, _DTypeT_co]: ... # Assignment preserves the receiver's shape; value broadcasting is gradual. def __setitem__(self, key: _ArrayIndex, value: ArrayLike, /) -> None: ... + # Unknown-rank arrays may yield scalars, so prefer the gradual overload. + @overload + def __iter__[N: IntVar](self: ndarray[[N], _DTypeT_co], /) -> Iterator[Any]: ... @overload def __iter__[N: IntVar, M: IntVar, Rest: _Shape = []]( self: ndarray[[N, M, *Rest], _DTypeT_co], / ) -> Iterator[ndarray[[M, *Rest], _DTypeT_co]]: ... - # Iteration over one-dimensional arrays yields NumPy scalars. - @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: ... diff --git a/tensor-shapes/pyrefly-numpy-stubs/test/test_indexing.py b/tensor-shapes/pyrefly-numpy-stubs/test/test_indexing.py index cf45fb9c89..53552444aa 100644 --- a/tensor-shapes/pyrefly-numpy-stubs/test/test_indexing.py +++ b/tensor-shapes/pyrefly-numpy-stubs/test/test_indexing.py @@ -39,6 +39,17 @@ def test_arange_optional_arguments() -> None: assert values.dtype == np.float32 +def test_iteration_preserves_unknown_rank() -> None: + def check_unknown(x: np.ndarray) -> None: + assert_type(next(iter(x)), Any) + + vector = np.ones(3) + matrix = np.ones((2, 3)) + check_unknown(vector) + check_unknown(matrix) + assert_shape(next(iter(matrix)).shape, (3,)) + + def test_array_protocol_shapes() -> None: vector = np.ones((5,)) matrix = np.ones((5, 3))