From ca4ac5bef18aaf78acaaab3e54417f85c4cd6a97 Mon Sep 17 00:00:00 2001 From: Brodi Elwood Date: Mon, 19 Oct 2026 19:42:34 -0400 Subject: [PATCH] Preserve unknown rank when iterating arrays --- tensor-shapes/jax-pyrefly-stubs/jax-stubs/_array.pyi | 4 ++-- tensor-shapes/jax-pyrefly-stubs/test/test_typing.py | 11 +++++++++++ .../pyrefly-numpy-stubs/numpy-stubs/__init__.pyi | 6 +++--- .../pyrefly-numpy-stubs/test/test_indexing.py | 11 +++++++++++ 4 files changed, 27 insertions(+), 5 deletions(-) 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 966a228bc3..87f9a7d578 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 7a32fce55e..394ab67bea 100644 --- a/tensor-shapes/pyrefly-numpy-stubs/numpy-stubs/__init__.pyi +++ b/tensor-shapes/pyrefly-numpy-stubs/numpy-stubs/__init__.pyi @@ -584,13 +584,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 __float__(self: ndarray[[]], /) -> float: ... def __complex__(self: ndarray[[]], /) -> complex: ... diff --git a/tensor-shapes/pyrefly-numpy-stubs/test/test_indexing.py b/tensor-shapes/pyrefly-numpy-stubs/test/test_indexing.py index 6acb7db7e0..9dde2b74a8 100644 --- a/tensor-shapes/pyrefly-numpy-stubs/test/test_indexing.py +++ b/tensor-shapes/pyrefly-numpy-stubs/test/test_indexing.py @@ -29,6 +29,17 @@ def test_arange_from_array_length() -> None: assert indices.dtype == np.dtype(np.intp) +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))