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

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

@@ -249,6 +249,27 @@
       "contributions": [
         "review"
       ]
+    },
+    {
+      "login": "paddyroddy",
+      "name": "Patrick J. Roddy",
+      "avatar_url": "https://avatars.githubusercontent.com/u/15052188?v=4",
+      "profile": "https://paddyroddy.github.io/",
+      "contributions": [
+        "code",
+        "doc",
+        "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.1 → 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.1 → array_api_extra-0.9.0}/CONTRIBUTORS.md

@@ -38,6 +38,8 @@ This project exists thanks to the following contributors
     <tr>
       <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.1 → array_api_extra-0.9.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: array-api-extra
-Version: 0.8.1
+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
@@ -139,6 +139,8 @@ This project exists thanks to the following contributors
     <tr>
       <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.1 → array_api_extra-0.9.0}/README.md

@@ -93,6 +93,8 @@ This project exists thanks to the following contributors
     <tr>
       <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.1 → array_api_extra-0.9.0}/docs/api-reference.md

@@ -16,6 +16,7 @@
     expand_dims
     isclose
     kron
+    nan_to_num
     nunique
     one_hot
     pad

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

@@ -58,7 +58,7 @@ array-api-extra = { path = ".", editable = true }
 typing-extensions = ">=4.14.1"
 pylint = ">=3.3.8"
 mypy = ">=1.17.1"
-basedpyright = ">=1.31.1"
+basedpyright = ">=1.31.3"
 numpydoc = ">=1.9.0,<2"
 # import dependencies for mypy:
 array-api-strict = ">=2.4.1,<2.5"
@@ -66,9 +66,9 @@ numpy = ">=2.1.3"
 hypothesis = ">=6.136.4"
 dask-core = ">=2025.7.0" # No distributed, tornado, etc.
 dprint = ">=0.50.0,<0.51"
-lefthook = ">=1.12.2,<2"
-ruff = ">=0.12.8,<0.13"
-typos = ">=1.35.3,<2"
+lefthook = ">=1.12.3,<2"
+ruff = ">=0.12.11,<0.13"
+typos = ">=1.35.5,<2"
 actionlint = ">=1.7.7,<2"
 blacken-docs = ">=1.19.1,<2"
 pytest = ">=8.4.1,<9"
@@ -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.1 → array_api_extra-0.9.0}/src/array_api_extra/__init__.py

@@ -1,6 +1,6 @@
 """Extra array functions built on top of the array API standard."""
 
-from ._delegation import isclose, one_hot, pad
+from ._delegation import isclose, nan_to_num, one_hot, pad
 from ._lib._at import at
 from ._lib._funcs import (
     apply_where,
@@ -17,7 +17,7 @@ from ._lib._funcs import (
 )
 from ._lib._lazy import lazy_apply
 
-__version__ = "0.8.1"
+__version__ = "0.9.0"
 
 # pylint: disable=duplicate-code
 __all__ = [
@@ -33,6 +33,7 @@ __all__ = [
     "isclose",
     "kron",
     "lazy_apply",
+    "nan_to_num",
     "nunique",
     "one_hot",
     "pad",

{array_api_extra-0.8.1 → array_api_extra-0.9.0}/src/array_api_extra/_delegation.py

@@ -18,7 +18,7 @@ from ._lib._utils._compat import device as get_device
 from ._lib._utils._helpers import asarrays
 from ._lib._utils._typing import Array, DType
 
-__all__ = ["isclose", "one_hot", "pad"]
+__all__ = ["isclose", "nan_to_num", "one_hot", "pad"]
 
 
 def isclose(
@@ -113,6 +113,85 @@ def isclose(
     return _funcs.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan, xp=xp)
 
 
+def nan_to_num(
+    x: Array | float | complex,
+    /,
+    *,
+    fill_value: int | float = 0.0,
+    xp: ModuleType | None = None,
+) -> Array:
+    """
+    Replace NaN with zero and infinity with large finite numbers (default behaviour).
+
+    If `x` is inexact, NaN is replaced by zero or by the user defined value in the
+    `fill_value` keyword, infinity is replaced by the largest finite floating
+    point value representable by ``x.dtype``, and -infinity is replaced by the
+    most negative finite floating point value representable by ``x.dtype``.
+
+    For complex dtypes, the above is applied to each of the real and
+    imaginary components of `x` separately.
+
+    Parameters
+    ----------
+    x : array | float | complex
+        Input data.
+    fill_value : int | float, optional
+        Value to be used to fill NaN values. If no value is passed
+        then NaN values will be replaced with 0.0.
+    xp : array_namespace, optional
+        The standard-compatible namespace for `x`. Default: infer.
+
+    Returns
+    -------
+    array
+        `x`, with the non-finite values replaced.
+
+    See Also
+    --------
+    array_api.isnan : Shows which elements are Not a Number (NaN).
+
+    Examples
+    --------
+    >>> import array_api_extra as xpx
+    >>> import array_api_strict as xp
+    >>> xpx.nan_to_num(xp.inf)
+    1.7976931348623157e+308
+    >>> xpx.nan_to_num(-xp.inf)
+    -1.7976931348623157e+308
+    >>> xpx.nan_to_num(xp.nan)
+    0.0
+    >>> x = xp.asarray([xp.inf, -xp.inf, xp.nan, -128, 128])
+    >>> xpx.nan_to_num(x)
+    array([ 1.79769313e+308, -1.79769313e+308,  0.00000000e+000, # may vary
+           -1.28000000e+002,  1.28000000e+002])
+    >>> y = xp.asarray([complex(xp.inf, xp.nan), xp.nan, complex(xp.nan, xp.inf)])
+    array([  1.79769313e+308,  -1.79769313e+308,   0.00000000e+000, # may vary
+         -1.28000000e+002,   1.28000000e+002])
+    >>> xpx.nan_to_num(y)
+    array([  1.79769313e+308 +0.00000000e+000j, # may vary
+             0.00000000e+000 +0.00000000e+000j,
+             0.00000000e+000 +1.79769313e+308j])
+    """
+    if isinstance(fill_value, complex):
+        msg = "Complex fill values are not supported."
+        raise TypeError(msg)
+
+    xp = array_namespace(x) if xp is None else xp
+
+    # for scalars we want to output an array
+    y = xp.asarray(x)
+
+    if (
+        is_cupy_namespace(xp)
+        or is_jax_namespace(xp)
+        or is_numpy_namespace(xp)
+        or is_torch_namespace(xp)
+    ):
+        return xp.nan_to_num(y, nan=fill_value)
+
+    return _funcs.nan_to_num(y, fill_value=fill_value, xp=xp)
+
+
 def one_hot(
     x: Array,
     /,

{array_api_extra-0.8.1 → array_api_extra-0.9.0}/src/array_api_extra/_lib/_funcs.py

@@ -738,6 +738,47 @@ def kron(
     return xp.reshape(result, res_shape)
 
 
+def nan_to_num(  # numpydoc ignore=PR01,RT01
+    x: Array,
+    /,
+    fill_value: int | float = 0.0,
+    *,
+    xp: ModuleType,
+) -> Array:
+    """See docstring in `array_api_extra._delegation.py`."""
+
+    def perform_replacements(  # numpydoc ignore=PR01,RT01
+        x: Array,
+        fill_value: int | float,
+        xp: ModuleType,
+    ) -> Array:
+        """Internal function to perform the replacements."""
+        x = xp.where(xp.isnan(x), fill_value, x)
+
+        # convert infinities to finite values
+        finfo = xp.finfo(x.dtype)
+        idx_posinf = xp.isinf(x) & ~xp.signbit(x)
+        idx_neginf = xp.isinf(x) & xp.signbit(x)
+        x = xp.where(idx_posinf, finfo.max, x)
+        return xp.where(idx_neginf, finfo.min, x)
+
+    if xp.isdtype(x.dtype, "complex floating"):
+        return perform_replacements(
+            xp.real(x),
+            fill_value,
+            xp,
+        ) + 1j * perform_replacements(
+            xp.imag(x),
+            fill_value,
+            xp,
+        )
+
+    if xp.isdtype(x.dtype, "numeric"):
+        return perform_replacements(x, fill_value, xp)
+
+    return x
+
+
 def nunique(x: Array, /, *, xp: ModuleType | None = None) -> Array:
     """
     Count the number of unique elements in an array.

{array_api_extra-0.8.1 → 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.1 → array_api_extra-0.9.0}/tests/conftest.py

@@ -232,3 +232,11 @@ def device(
     if library == Backend.TORCH_GPU:
         return xp.device("cpu")
     return get_device(xp.empty(0))
+
+
+@pytest.fixture
+def infinity(library: Backend) -> float:
+    """Retrieve the positive infinity value for the given backend."""
+    if library in (Backend.TORCH, Backend.TORCH_GPU):
+        return 3.4028235e38
+    return 1.7976931348623157e308

{array_api_extra-0.8.1 → array_api_extra-0.9.0}/tests/test_funcs.py

@@ -21,6 +21,7 @@ from array_api_extra import (
     expand_dims,
     isclose,
     kron,
+    nan_to_num,
     nunique,
     one_hot,
     pad,
@@ -40,6 +41,7 @@ lazy_xp_function(cov)
 lazy_xp_function(create_diagonal)
 lazy_xp_function(expand_dims)
 lazy_xp_function(kron)
+lazy_xp_function(nan_to_num)
 lazy_xp_function(nunique)
 lazy_xp_function(one_hot)
 lazy_xp_function(pad)
@@ -941,6 +943,140 @@ class TestKron:
         xp_assert_equal(kron(a, b, xp=xp), k)
 
 
+class TestNanToNum:
+    def test_bool(self, xp: ModuleType) -> None:
+        a = xp.asarray([True])
+        xp_assert_equal(nan_to_num(a, xp=xp), a)
+
+    def test_scalar_pos_inf(self, xp: ModuleType, infinity: float) -> None:
+        a = xp.inf
+        xp_assert_equal(nan_to_num(a, xp=xp), xp.asarray(infinity))
+
+    def test_scalar_neg_inf(self, xp: ModuleType, infinity: float) -> None:
+        a = -xp.inf
+        xp_assert_equal(nan_to_num(a, xp=xp), -xp.asarray(infinity))
+
+    def test_scalar_nan(self, xp: ModuleType) -> None:
+        a = xp.nan
+        xp_assert_equal(nan_to_num(a, xp=xp), xp.asarray(0.0))
+
+    def test_real(self, xp: ModuleType, infinity: float) -> None:
+        a = xp.asarray([xp.inf, -xp.inf, xp.nan, -128, 128])
+        xp_assert_equal(
+            nan_to_num(a, xp=xp),
+            xp.asarray(
+                [
+                    infinity,
+                    -infinity,
+                    0.0,
+                    -128,
+                    128,
+                ]
+            ),
+        )
+
+    def test_complex(self, xp: ModuleType, infinity: float) -> None:
+        a = xp.asarray(
+            [
+                complex(xp.inf, xp.nan),
+                xp.nan,
+                complex(xp.nan, xp.inf),
+            ]
+        )
+        xp_assert_equal(
+            nan_to_num(a),
+            xp.asarray([complex(infinity, 0), complex(0, 0), complex(0, infinity)]),
+        )
+
+    def test_empty_array(self, xp: ModuleType) -> None:
+        a = xp.asarray([], dtype=xp.float32)  # forced dtype due to torch
+        xp_assert_equal(nan_to_num(a, xp=xp), a)
+        assert xp.isdtype(nan_to_num(a, xp=xp).dtype, xp.float32)
+
+    @pytest.mark.parametrize(
+        ("in_vals", "fill_value", "out_vals"),
+        [
+            ([1, 2, np.nan, 4], 3, [1.0, 2.0, 3.0, 4.0]),
+            ([1, 2, np.nan, 4], 3.0, [1.0, 2.0, 3.0, 4.0]),
+            (
+                [
+                    complex(1, 1),
+                    complex(2, 2),
+                    complex(np.nan, 0),
+                    complex(4, 4),
+                ],
+                3,
+                [
+                    complex(1.0, 1.0),
+                    complex(2.0, 2.0),
+                    complex(3.0, 0.0),
+                    complex(4.0, 4.0),
+                ],
+            ),
+            (
+                [
+                    complex(1, 1),
+                    complex(2, 2),
+                    complex(0, np.nan),
+                    complex(4, 4),
+                ],
+                3.0,
+                [
+                    complex(1.0, 1.0),
+                    complex(2.0, 2.0),
+                    complex(0.0, 3.0),
+                    complex(4.0, 4.0),
+                ],
+            ),
+            (
+                [
+                    complex(1, 1),
+                    complex(2, 2),
+                    complex(np.nan, np.nan),
+                    complex(4, 4),
+                ],
+                3.0,
+                [
+                    complex(1.0, 1.0),
+                    complex(2.0, 2.0),
+                    complex(3.0, 3.0),
+                    complex(4.0, 4.0),
+                ],
+            ),
+        ],
+    )
+    def test_fill_value_success(
+        self,
+        xp: ModuleType,
+        in_vals: Array,
+        fill_value: int | float,
+        out_vals: Array,
+    ) -> None:
+        a = xp.asarray(in_vals)
+        xp_assert_equal(
+            nan_to_num(a, fill_value=fill_value, xp=xp),
+            xp.asarray(out_vals),
+        )
+
+    def test_fill_value_failure(self, xp: ModuleType) -> None:
+        a = xp.asarray(
+            [
+                complex(1, 1),
+                complex(xp.nan, xp.nan),
+                complex(3, 3),
+            ]
+        )
+        with pytest.raises(
+            TypeError,
+            match="Complex fill values are not supported",
+        ):
+            _ = nan_to_num(
+                a,
+                fill_value=complex(2, 2),  # type: ignore[arg-type] # pyright: ignore[reportArgumentType]
+                xp=xp,
+            )
+
+
 class TestNUnique:
     def test_simple(self, xp: ModuleType):
         a = xp.asarray([[1, 1], [0, 2], [2, 2]])

{array_api_extra-0.8.1 → 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.1 → 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)