Preserve unknown rank when iterating arrays

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: ...

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`
 

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: ...

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