replace

docs/sphinx/api-manipulation.md

@@ -10,4 +10,5 @@
     broadcast_shapes
     expand_dims
     pad
+    replace
 ```

src/array_api_extra/__init__.py

@@ -9,7 +9,7 @@
 from ._indexing import diag_indices, tril_indices, triu_indices, unravel_index
 from ._lazy import lazy_apply
 from ._linalg import kron
-from ._manipulation import atleast_nd, broadcast_shapes, expand_dims, pad
+from ._manipulation import atleast_nd, broadcast_shapes, expand_dims, pad, replace
 from ._searching import searchsorted
 from ._set import isin, nunique, setdiff1d, union1d
 from ._sorting import argpartition, partition
@@ -45,6 +45,7 @@
     "pad",
     "partition",
     "rad2deg",
+    "replace",
     "searchsorted",
     "setdiff1d",
     "sinc",

src/array_api_extra/_agnostic/_manipulation.py

@@ -1,5 +1,6 @@
 """Array-agnostic implementations for manipulation functions."""
 
+import cmath
 import math
 import typing
 from collections.abc import Sequence
@@ -8,7 +9,7 @@
 from .._lib import _compat, _helpers
 from .._lib._typing import Array, ArrayNamespace
 
-__all__ = ["atleast_nd", "broadcast_shapes", "expand_dims", "pad"]
+__all__ = ["atleast_nd", "broadcast_shapes", "expand_dims", "pad", "replace"]
 
 
 def atleast_nd(x: Array, /, *, ndim: int, xp: ArrayNamespace) -> Array:
@@ -97,3 +98,11 @@ def pad(
         device=_compat.device(x),
     )
     return at(padded, tuple(slices)).set(x)
+
+
+def replace(  # numpydoc ignore=PR01,RT01
+    x: Array, old: complex, new: complex, /, *, xp: ArrayNamespace
+) -> Array:
+    """See docstring in `array_api_extra._manipulation`."""
+    mask = xp.isnan(x) if cmath.isnan(old) else x == old
+    return xp.where(mask, new, x)

src/array_api_extra/_manipulation.py

@@ -7,7 +7,7 @@
 from ._lib import _compat, _helpers
 from ._lib._typing import Array, ArrayNamespace
 
-__all__ = ["atleast_nd", "broadcast_shapes", "expand_dims", "pad"]
+__all__ = ["atleast_nd", "broadcast_shapes", "expand_dims", "pad", "replace"]
 
 
 def atleast_nd(x: Array, /, *, ndim: int, xp: ArrayNamespace | None = None) -> Array:
@@ -282,3 +282,72 @@ def pad(
     return _agnostic._manipulation.pad(
         x, pad_width, constant_values=constant_values, xp=xp
     )
+
+
+def replace(
+    x: Array,
+    old: complex,
+    new: complex,
+    /,
+    *,
+    xp: ArrayNamespace | None = None,
+) -> Array:
+    """
+    Replace each element of an array equal to a scalar with another scalar.
+
+    Parameters
+    ----------
+    x : array
+        Input array.
+    old : python scalar
+        Value to replace. If NaN, every NaN element of `x` is replaced.
+    new : python scalar
+        Replacement value.
+    xp : array_namespace, optional
+        The standard-compatible namespace for `x`. Default: infer.
+
+    Returns
+    -------
+    array
+        An array with the same shape and data type as `x`, in which each element
+        equal to `old` is replaced by `new` and all other elements are equal to the
+        corresponding elements of `x`. `x` is not modified.
+
+    Notes
+    -----
+    `old` and `new` should have types and values compatible with the data type of
+    `x` (see :ref:`array-api:mixing-scalars-and-arrays`). Otherwise, the behaviour
+    depends on the backend.
+
+    For sparse arrays, elements which are not explicitly stored are replaced too:
+    if `old` equals the fill value, the fill value of the result is `new`.
+    The result keeps the sparse format of `x` and is never densified.
+
+    Examples
+    --------
+    >>> import array_api_strict as xp
+    >>> import array_api_extra as xpx
+    >>> x = xp.asarray([1, 2, 1, 3])
+    >>> xpx.replace(x, 1, 0)
+    Array([0, 2, 0, 3], dtype=array_api_strict.int64)
+
+    NaN elements are replaced when `old` is NaN:
+
+    >>> x = xp.asarray([1.0, xp.nan, 3.0])
+    >>> xpx.replace(x, xp.nan, 0.0)
+    Array([1., 0., 3.], dtype=array_api_strict.float64)
+    """
+    xp = _compat.array_namespace(x) if xp is None else xp
+
+    if _compat.is_pydata_sparse_namespace(xp):
+        # `elemwise` applies the replacement to both the stored elements and the
+        # fill value, so implicit elements are replaced without densifying
+        def replace_dense(data: Array) -> Array:  # numpydoc ignore=PR01,RT01
+            """Replace in the dense arrays underlying a sparse array."""
+            return _agnostic._manipulation.replace(
+                data, old, new, xp=_compat.array_namespace(data)
+            )
+
+        return xp.elemwise(replace_dense, x)
+
+    return _agnostic._manipulation.replace(x, old, new, xp=xp)

tests/main/test_manipulation.py

@@ -1,9 +1,10 @@
 import math
+from typing import Any, cast
 
 import numpy as np
 import pytest
 
-from array_api_extra import atleast_nd, broadcast_shapes, expand_dims, pad
+from array_api_extra import atleast_nd, broadcast_shapes, expand_dims, pad, replace
 from array_api_extra._lib._backends import Backend
 from array_api_extra._lib._compat import device as get_device
 from array_api_extra._lib._typing import ArrayNamespace, Device
@@ -13,6 +14,7 @@
 lazy_xp_function(broadcast_shapes)
 lazy_xp_function(expand_dims)
 lazy_xp_function(pad)
+lazy_xp_function(replace)
 
 
 class TestAtLeastND:
@@ -371,3 +373,86 @@ def test_sequence_of_tuples_width(self, xp: ArrayNamespace):
         assert padded.shape == (4, 6)
         padded = pad(a, ((1, 0), (0, 0)))
         assert padded.shape == (4, 4)
+
+
+class TestReplace:
+    def test_simple(self, xp: ArrayNamespace):
+        x = xp.asarray([1, 2, 1, 3])
+        assert_equal(replace(x, 1, 5), xp.asarray([5, 2, 5, 3]))
+
+    def test_no_match(self, xp: ArrayNamespace):
+        x = xp.asarray([1, 2, 3])
+        assert_equal(replace(x, 4, 5), x)
+
+    def test_zero(self, xp: ArrayNamespace):
+        # Zero is the fill value of sparse arrays
+        x = xp.asarray([0, 1, 0, 2])
+        assert_equal(replace(x, 0, 5), xp.asarray([5, 1, 5, 2]))
+        assert_equal(replace(x, 1, 0), xp.asarray([0, 0, 0, 2]))
+
+    def test_nan(self, xp: ArrayNamespace):
+        x = xp.asarray([1.0, xp.nan, 2.0, xp.nan])
+        assert_equal(replace(x, xp.nan, 0.0), xp.asarray([1.0, 0.0, 2.0, 0.0]))
+        assert_equal(replace(x, 1.0, xp.nan), xp.asarray([xp.nan, xp.nan, 2.0, xp.nan]))
+
+    def test_complex(self, xp: ArrayNamespace):
+        x = xp.asarray([1 + 1j, 2 + 0j, complex(xp.nan, 1)])
+        expect = xp.asarray([0j, 2 + 0j, complex(xp.nan, 1)])
+        assert_equal(replace(x, 1 + 1j, 0j), expect)
+        # Complex elements with a NaN component are NaN
+        expect = xp.asarray([1 + 1j, 2 + 0j, 3j])
+        assert_equal(replace(x, complex(xp.nan, 0), 3j), expect)
+
+    def test_bool(self, xp: ArrayNamespace):
+        x = xp.asarray([True, False, True])
+        assert_equal(replace(x, True, False), xp.asarray([False, False, False]))
+
+    @pytest.mark.parametrize("dtype", ["int8", "uint8", "int32", "float32", "float64"])
+    def test_dtype(self, xp: ArrayNamespace, dtype: str):
+        x = xp.asarray([1, 2, 1], dtype=getattr(xp, dtype))
+        expect = xp.asarray([3, 2, 3], dtype=getattr(xp, dtype))
+        assert_equal(replace(x, 1, 3), expect)
+
+    def test_0d(self, xp: ArrayNamespace):
+        assert_equal(replace(xp.asarray(1.0), 1.0, 2.0), xp.asarray(2.0))
+        assert_equal(replace(xp.asarray(1.0), 3.0, 2.0), xp.asarray(1.0))
+
+    def test_ndim(self, xp: ArrayNamespace):
+        x = xp.asarray(np.reshape(np.arange(2 * 3 * 4) % 3, (2, 3, 4)))
+        expect = xp.asarray(np.reshape(np.arange(2 * 3 * 4) % 3 * 4 % 5, (2, 3, 4)))
+        actual = replace(x, 1, 4)
+        assert_equal(replace(actual, 2, 3), expect)
+
+    def test_input_not_modified(self, xp: ArrayNamespace):
+        x = xp.asarray([1, 2, 1])
+        _ = replace(x, 1, 5)
+        assert_equal(x, xp.asarray([1, 2, 1]))
+
+    def test_device(self, xp: ArrayNamespace, device: Device):
+        x = xp.asarray([1.0, 2.0], device=device)
+        assert get_device(replace(x, 1.0, 3.0)) == device
+
+    def test_xp(self, xp: ArrayNamespace):
+        x = xp.asarray([1, 2])
+        assert_equal(replace(x, 1, 3, xp=xp), xp.asarray([3, 2]))
+
+    @pytest.mark.parametrize("fmt", ["coo", "gcxs", "dok"])
+    def test_sparse(self, fmt: str):
+        sparse = pytest.importorskip("sparse")
+        x: Any = sparse.asarray(np.asarray([0.0, 1.0, 0.0, np.nan]), format=fmt)
+
+        # Replacing the fill value changes the fill value, without densifying
+        actual = cast(Any, replace(x, 0.0, 5.0))
+        assert type(actual) is type(x)
+        assert actual.fill_value == 5.0
+        assert actual.nnz == x.nnz
+        np.testing.assert_array_equal(actual.todense(), [5.0, 1.0, 5.0, np.nan])
+
+        actual = cast(Any, replace(x, np.nan, 2.0))
+        assert type(actual) is type(x)
+        assert actual.fill_value == 0.0
+        np.testing.assert_array_equal(actual.todense(), [0.0, 1.0, 0.0, 2.0])
+
+        # Input is not modified
+        assert x.fill_value == 0.0
+        np.testing.assert_array_equal(x.todense(), [0.0, 1.0, 0.0, np.nan])