array-api-extra 0.8.2__tar.gz → 0.9.0__tar.gz

{array_api_extra-0.8.2 → array_api_extra-0.9.0}/.all-contributorsrc

@@ -261,6 +261,15 @@
         "example",
         "test"
       ]
+    },
+    {
+      "login": "amacati",
+      "name": "Martin Schuck",
+      "avatar_url": "https://avatars.githubusercontent.com/u/57562633?v=4",
+      "profile": "https://amacati.github.io/",
+      "contributions": [
+        "ideas"
+      ]
     }
   ]
 }

{array_api_extra-0.8.2 → array_api_extra-0.9.0}/.github/workflows/ci.yml

@@ -48,6 +48,7 @@ jobs:
           - tests-py313
           - tests-numpy1
           - tests-backends
+          - tests-backends-py310
           - tests-nogil
         runs-on: [ubuntu-latest]
 

{array_api_extra-0.8.2 → array_api_extra-0.9.0}/CONTRIBUTORS.md

@@ -39,6 +39,7 @@ This project exists thanks to the following contributors
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/lithomas1"><img src="https://avatars.githubusercontent.com/u/47963215?v=4?s=100" width="100px;" alt="Thomas Li"/><br /><sub><b>Thomas Li</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/issues?q=author%3Alithomas1" title="Bug reports">🐛</a> <a href="#tool-lithomas1" title="Tools">🔧</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/pearu"><img src="https://avatars.githubusercontent.com/u/402156?v=4?s=100" width="100px;" alt="Pearu Peterson"/><br /><sub><b>Pearu Peterson</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/pulls?q=is%3Apr+reviewed-by%3Apearu" title="Reviewed Pull Requests">👀</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://paddyroddy.github.io/"><img src="https://avatars.githubusercontent.com/u/15052188?v=4?s=100" width="100px;" alt="Patrick J. Roddy"/><br /><sub><b>Patrick J. Roddy</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/commits?author=paddyroddy" title="Code">💻</a> <a href="https://github.com/data-apis/array-api-extra/commits?author=paddyroddy" title="Documentation">📖</a> <a href="#example-paddyroddy" title="Examples">💡</a> <a href="https://github.com/data-apis/array-api-extra/commits?author=paddyroddy" title="Tests">⚠️</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://amacati.github.io/"><img src="https://avatars.githubusercontent.com/u/57562633?v=4?s=100" width="100px;" alt="Martin Schuck"/><br /><sub><b>Martin Schuck</b></sub></a><br /><a href="#ideas-amacati" title="Ideas, Planning, & Feedback">🤔</a></td>
     </tr>
   </tbody>
 </table>

{array_api_extra-0.8.2 → array_api_extra-0.9.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: array-api-extra
-Version: 0.8.2
+Version: 0.9.0
 Summary: Extra array functions built on top of the array API standard.
 Project-URL: Homepage, https://github.com/data-apis/array-api-extra
 Project-URL: Bug Tracker, https://github.com/data-apis/array-api-extra/issues
@@ -140,6 +140,7 @@ This project exists thanks to the following contributors
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/lithomas1"><img src="https://avatars.githubusercontent.com/u/47963215?v=4?s=100" width="100px;" alt="Thomas Li"/><br /><sub><b>Thomas Li</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/issues?q=author%3Alithomas1" title="Bug reports">🐛</a> <a href="#tool-lithomas1" title="Tools">🔧</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/pearu"><img src="https://avatars.githubusercontent.com/u/402156?v=4?s=100" width="100px;" alt="Pearu Peterson"/><br /><sub><b>Pearu Peterson</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/pulls?q=is%3Apr+reviewed-by%3Apearu" title="Reviewed Pull Requests">👀</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://paddyroddy.github.io/"><img src="https://avatars.githubusercontent.com/u/15052188?v=4?s=100" width="100px;" alt="Patrick J. Roddy"/><br /><sub><b>Patrick J. Roddy</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/commits?author=paddyroddy" title="Code">💻</a> <a href="https://github.com/data-apis/array-api-extra/commits?author=paddyroddy" title="Documentation">📖</a> <a href="#example-paddyroddy" title="Examples">💡</a> <a href="https://github.com/data-apis/array-api-extra/commits?author=paddyroddy" title="Tests">⚠️</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://amacati.github.io/"><img src="https://avatars.githubusercontent.com/u/57562633?v=4?s=100" width="100px;" alt="Martin Schuck"/><br /><sub><b>Martin Schuck</b></sub></a><br /><a href="#ideas-amacati" title="Ideas, Planning, & Feedback">🤔</a></td>
     </tr>
   </tbody>
 </table>

{array_api_extra-0.8.2 → array_api_extra-0.9.0}/README.md

@@ -94,6 +94,7 @@ This project exists thanks to the following contributors
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/lithomas1"><img src="https://avatars.githubusercontent.com/u/47963215?v=4?s=100" width="100px;" alt="Thomas Li"/><br /><sub><b>Thomas Li</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/issues?q=author%3Alithomas1" title="Bug reports">🐛</a> <a href="#tool-lithomas1" title="Tools">🔧</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/pearu"><img src="https://avatars.githubusercontent.com/u/402156?v=4?s=100" width="100px;" alt="Pearu Peterson"/><br /><sub><b>Pearu Peterson</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/pulls?q=is%3Apr+reviewed-by%3Apearu" title="Reviewed Pull Requests">👀</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://paddyroddy.github.io/"><img src="https://avatars.githubusercontent.com/u/15052188?v=4?s=100" width="100px;" alt="Patrick J. Roddy"/><br /><sub><b>Patrick J. Roddy</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/commits?author=paddyroddy" title="Code">💻</a> <a href="https://github.com/data-apis/array-api-extra/commits?author=paddyroddy" title="Documentation">📖</a> <a href="#example-paddyroddy" title="Examples">💡</a> <a href="https://github.com/data-apis/array-api-extra/commits?author=paddyroddy" title="Tests">⚠️</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://amacati.github.io/"><img src="https://avatars.githubusercontent.com/u/57562633?v=4?s=100" width="100px;" alt="Martin Schuck"/><br /><sub><b>Martin Schuck</b></sub></a><br /><a href="#ideas-amacati" title="Ideas, Planning, & Feedback">🤔</a></td>
     </tr>
   </tbody>
 </table>

{array_api_extra-0.8.2 → array_api_extra-0.9.0}/pyproject.toml

@@ -67,7 +67,7 @@ hypothesis = ">=6.136.4"
 dask-core = ">=2025.7.0" # No distributed, tornado, etc.
 dprint = ">=0.50.0,<0.51"
 lefthook = ">=1.12.3,<2"
-ruff = ">=0.12.8,<0.13"
+ruff = ">=0.12.11,<0.13"
 typos = ">=1.35.5,<2"
 actionlint = ">=1.7.7,<2"
 blacken-docs = ">=1.19.1,<2"
@@ -155,13 +155,16 @@ dask-core = ">=2025.7.0" # No distributed, tornado, etc.
 sparse = ">=0.17.0"
 
 [tool.pixi.feature.backends.target.linux-64.dependencies]
-jax = ">=0.6.0,!=0.6.2" # 0.6.2 segfaults on Linux CUDA
+# On CPU Python 3.10, use 0.6.2
+# On CPU Python >=3.11, use >=0.7.0
+# On GPU, use 0.6.0 (0.6.2 and 0.7.0 both segfault); see jaxlib pin below.
+jax = ">=0.6.0"
 
 [tool.pixi.feature.backends.target.osx-64.dependencies]
-jax = ">=0.6.0,!=0.6.2"
+jax = ">=0.6.0"
 
 [tool.pixi.feature.backends.target.osx-arm64.dependencies]
-jax = ">=0.6.0,!=0.6.2"
+jax = ">=0.6.0"
 
 [tool.pixi.feature.backends.target.win-64.dependencies]
 # jax = "*"  # unavailable
@@ -176,8 +179,9 @@ jax = ">=0.6.0,!=0.6.2"
 system-requirements = { cuda = "12" }
 
 [tool.pixi.feature.cuda-backends.target.linux-64.dependencies]
-cupy = ">=13.5.1"
-jaxlib = { version = ">=0.6.0", build = "cuda12*" }
+cupy = ">=13.6.0"
+# JAX 0.6.2 and 0.7.0 segfault on CUDA
+jaxlib = { version = ">=0.6.0,!=0.6.2,!=0.7.0", build = "cuda12*" }
 pytorch = { version = ">=2.7.1", build = "cuda12*" }
 
 [tool.pixi.feature.cuda-backends.target.osx-64.dependencies]
@@ -191,13 +195,13 @@ pytorch = { version = ">=2.7.1", build = "cuda12*" }
 # pytorch = { version = "*", build = "cuda12*" }  # unavailable
 
 [tool.pixi.feature.cuda-backends.target.win-64.dependencies]
-cupy = ">=13.5.1"
+cupy = ">=13.6.0"
 # jaxlib = { version = "*", build = "cuda12*" }  # unavailable
 pytorch = { version = ">=2.7.1", build = "cuda12*" }
 
 [tool.pixi.feature.nogil.dependencies]
 python-freethreading = "~=3.13.0"
-pytest-run-parallel = ">=0.6.0"
+pytest-run-parallel = ">=0.6.1"
 numpy = ">=2.3.2"
 # pytorch = "*"  # Not available on Python 3.13t yet
 dask-core = ">=2025.7.0" # No distributed, tornado, etc.
@@ -212,12 +216,16 @@ tests = { features = ["py313", "tests"], solve-group = "py313" }
 tests-py313 = { features = ["py313", "tests"], solve-group = "py313" } # alias of tests
 
 # Some backends may pin numpy; use separate solve-group
-dev = { features = ["py310", "lint", "tests", "docs", "dev", "backends"], solve-group = "backends" }
-tests-backends = { features = ["py310", "tests", "backends"], solve-group = "backends" }
+dev = { features = ["py313", "lint", "tests", "docs", "dev", "backends"], solve-group = "backends" }
+tests-backends = { features = ["py313", "tests", "backends"], solve-group = "backends" }
+# Note: Python 3.10 has already been dropped by some backends (like JAX),
+# so this is testing older versions.
+tests-backends-py310 = { features = ["py310", "tests", "backends"] }
 
 # CUDA not available on free github actions and on some developers' PCs
-dev-cuda = { features = ["py310", "lint", "tests", "docs", "dev", "backends", "cuda-backends"], solve-group = "cuda" }
-tests-cuda = { features = ["py310", "tests", "backends", "cuda-backends"], solve-group = "cuda" }
+dev-cuda = { features = ["py313", "lint", "tests", "docs", "dev", "backends", "cuda-backends"], solve-group = "cuda" }
+tests-cuda = { features = ["py313", "tests", "backends", "cuda-backends"], solve-group = "cuda" }
+tests-cuda-py310 = { features = ["py310", "tests", "backends", "cuda-backends"] }
 
 # Ungrouped environments
 tests-numpy1 = ["py310", "tests", "numpy1"]

{array_api_extra-0.8.2 → array_api_extra-0.9.0}/src/array_api_extra/__init__.py

@@ -17,7 +17,7 @@ from ._lib._funcs import (
 )
 from ._lib._lazy import lazy_apply
 
-__version__ = "0.8.2"
+__version__ = "0.9.0"
 
 # pylint: disable=duplicate-code
 __all__ = [

{array_api_extra-0.8.2 → array_api_extra-0.9.0}/src/array_api_extra/_lib/_utils/_helpers.py

@@ -6,7 +6,7 @@ import io
 import math
 import pickle
 import types
-from collections.abc import Callable, Generator, Iterable
+from collections.abc import Callable, Generator, Iterable, Iterator
 from functools import wraps
 from types import ModuleType
 from typing import (
@@ -512,13 +512,24 @@ class _AutoJITWrapper(Generic[T]):  # numpydoc ignore=PR01
     convert them to/from PyTrees.
     """
 
-    obj: T
+    _obj: Any
+    _is_iter: bool
     _registered: ClassVar[bool] = False
-    __slots__: tuple[str, ...] = ("obj",)
+    __slots__: tuple[str, ...] = ("_is_iter", "_obj")
 
     def __init__(self, obj: T) -> None:  # numpydoc ignore=GL08
         self._register()
-        self.obj = obj
+        if isinstance(obj, Iterator):
+            self._obj = list(obj)
+            self._is_iter = True
+        else:
+            self._obj = obj
+            self._is_iter = False
+
+    @property
+    def obj(self) -> T:  # numpydoc ignore=RT01
+        """Return wrapped object."""
+        return iter(self._obj) if self._is_iter else self._obj
 
     @classmethod
     def _register(cls) -> None:  # numpydoc ignore=SS06
@@ -531,7 +542,7 @@ class _AutoJITWrapper(Generic[T]):  # numpydoc ignore=PR01
 
             jax.tree_util.register_pytree_node(
                 cls,
-                lambda obj: pickle_flatten(obj, jax.Array),  # pyright: ignore[reportUnknownArgumentType]
+                lambda instance: pickle_flatten(instance, jax.Array),  # pyright: ignore[reportUnknownArgumentType]
                 lambda aux_data, children: pickle_unflatten(children, aux_data),  # pyright: ignore[reportUnknownArgumentType]
             )
             cls._registered = True
@@ -556,6 +567,7 @@ def jax_autojit(
     - Automatically descend into non-array return values and find ``jax.Array`` objects
       inside them, then rebuild them downstream of exiting the JIT, swapping the JAX
       tracer objects with concrete arrays.
+    - Returned iterators are immediately completely consumed.
 
     See Also
     --------

{array_api_extra-0.8.2 → array_api_extra-0.9.0}/tests/test_helpers.py

@@ -1,3 +1,4 @@
+from collections.abc import Iterator
 from types import ModuleType
 from typing import TYPE_CHECKING, Generic, TypeVar, cast
 
@@ -417,3 +418,16 @@ class TestJAXAutoJIT:
         out = f([1, 2])
         assert isinstance(out, list)
         assert out == [3, 4]
+
+    def test_iterators(self, jnp: ModuleType):
+        @jax_autojit
+        def f(x: Array) -> Iterator[Array]:
+            return (x + i for i in range(2))
+
+        inp = jnp.asarray([1, 2])
+        out = f(inp)
+        assert isinstance(out, Iterator)
+        xp_assert_equal(next(out), jnp.asarray([1, 2]))
+        xp_assert_equal(next(out), jnp.asarray([2, 3]))
+        with pytest.raises(StopIteration):
+            _ = next(out)

{array_api_extra-0.8.2 → array_api_extra-0.9.0}/tests/test_testing.py

@@ -1,4 +1,4 @@
-from collections.abc import Callable
+from collections.abc import Callable, Iterator
 from types import ModuleType
 from typing import cast
 
@@ -468,3 +468,22 @@ def test_patch_lazy_xp_functions_deprecated_monkeypatch(
     monkeypatch.undo()
     y = non_materializable5(x)
     xp_assert_equal(y, x)
+
+
+def my_iter(x: Array) -> Iterator[Array]:
+    yield x[0, :]
+    yield x[1, :]
+
+
+lazy_xp_function(my_iter)
+
+
+def test_patch_lazy_xp_functions_iter(xp: ModuleType):
+    x = xp.asarray([[1.0, 2.0], [3.0, 4.0]])
+    it = my_iter(x)
+
+    assert isinstance(it, Iterator)
+    xp_assert_equal(next(it), x[0, :])
+    xp_assert_equal(next(it), x[1, :])
+    with pytest.raises(StopIteration):
+        _ = next(it)