From 8ddf4bcc88afd06d907f630cc852a66d41b12152 Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Thu, 29 Oct 2026 14:25:50 -0400 Subject: [PATCH 1/2] first draft --- src/array_api_extra/_manipulation.py | 49 ++++++++++++++++++++++++++++ 1 file changed, 49 insertions(+) diff --git a/src/array_api_extra/_manipulation.py b/src/array_api_extra/_manipulation.py index fa3eae9b..3ca73945 100644 --- a/src/array_api_extra/_manipulation.py +++ b/src/array_api_extra/_manipulation.py @@ -282,3 +282,52 @@ def pad( return _agnostic._manipulation.pad( x, pad_width, constant_values=constant_values, xp=xp ) + +def replace( + x: array, + old: Union[bool, int, float, complex], + new: Union[bool, int, float, complex], + /, +) -> array: + """ + Replaces each element of an input array ``x`` which is equal to a scalar value ``old`` with a scalar value ``new``. + + Parameters + ---------- + x: array + input array. + old: Union[bool, int, float, complex] + scalar value to replace. **Should** have a type and value compatible with the data type of ``x`` (see :ref:`mixing-scalars-and-arrays`). + new: Union[bool, int, float, complex] + replacement scalar value. **Should** have a type and value compatible with the data type of ``x`` (see :ref:`mixing-scalars-and-arrays`). + copy: Union[bool, None] + return whether the operation should return a copy + + Returns + ------- + out: array + an array in which each element ``x_i`` equal to ``old`` is replaced by ``new`` and all other elements are equal to the corresponding elements in ``x``. The returned array **must** have the same shape and data type as ``x``. + + Notes + ----- + + - An element ``x_i`` **must** be replaced if ``x_i`` is equal to ``old`` (see :func:`~array_api.equal`) or if both ``x_i`` and ``old`` are ``NaN`` (see :func:`~array_api.isnan`). + - The function **must** consider every element of ``x``, including elements which are not explicitly stored (e.g., implicit fill values in a sparse array representation). + - The function **must not** modify ``x``. + - If ``old`` or ``new`` does not have a type and value compatible with the data type of ``x``, behavior is unspecified and thus implementation-defined. + """ + if _compat.is_pydata_sparse_namespace(xp): + if isinstance(arr, sp.DOK): + if copy == False: + throw(ArgumentError("cannot zero-copy change fill value of DOK format")) + arr = arr.asformat("coo") + if copy == True: + arr = copy(arr) + if arr.fill_value != old: + throw(densification warning) + replace(arr.val, copy=false) + return arr.to_dense + else: + replace(arr.val, copy=false) + arr.fill_value = new + return _agnostic.where(_agnostic.isnan(arr) if old != old else arr == old, new, arr) \ No newline at end of file From 4da058ef8ebd6cc2fabcc010a5dce05eac49d510 Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Thu, 29 Oct 2026 14:41:50 -0400 Subject: [PATCH 2/2] replace --- docs/sphinx/api-manipulation.md | 1 + src/array_api_extra/__init__.py | 3 +- .../_agnostic/_manipulation.py | 11 ++- src/array_api_extra/_manipulation.py | 88 ++++++++++++------- tests/main/test_manipulation.py | 87 +++++++++++++++++- 5 files changed, 153 insertions(+), 37 deletions(-) diff --git a/docs/sphinx/api-manipulation.md b/docs/sphinx/api-manipulation.md index b8875fd5..a778786b 100644 --- a/docs/sphinx/api-manipulation.md +++ b/docs/sphinx/api-manipulation.md @@ -10,4 +10,5 @@ broadcast_shapes expand_dims pad + replace ``` diff --git a/src/array_api_extra/__init__.py b/src/array_api_extra/__init__.py index c4e5d42a..d108ca4d 100644 --- a/src/array_api_extra/__init__.py +++ b/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", diff --git a/src/array_api_extra/_agnostic/_manipulation.py b/src/array_api_extra/_agnostic/_manipulation.py index fb28faac..0d2770ef 100644 --- a/src/array_api_extra/_agnostic/_manipulation.py +++ b/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) diff --git a/src/array_api_extra/_manipulation.py b/src/array_api_extra/_manipulation.py index 3ca73945..9bd2c071 100644 --- a/src/array_api_extra/_manipulation.py +++ b/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: @@ -283,51 +283,71 @@ def pad( x, pad_width, constant_values=constant_values, xp=xp ) + def replace( - x: array, - old: Union[bool, int, float, complex], - new: Union[bool, int, float, complex], + x: Array, + old: complex, + new: complex, /, -) -> array: + *, + xp: ArrayNamespace | None = None, +) -> Array: """ - Replaces each element of an input array ``x`` which is equal to a scalar value ``old`` with a scalar value ``new``. + Replace each element of an array equal to a scalar with another scalar. Parameters ---------- - x: array - input array. - old: Union[bool, int, float, complex] - scalar value to replace. **Should** have a type and value compatible with the data type of ``x`` (see :ref:`mixing-scalars-and-arrays`). - new: Union[bool, int, float, complex] - replacement scalar value. **Should** have a type and value compatible with the data type of ``x`` (see :ref:`mixing-scalars-and-arrays`). - copy: Union[bool, None] - return whether the operation should return a copy + 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 ------- - out: array - an array in which each element ``x_i`` equal to ``old`` is replaced by ``new`` and all other elements are equal to the corresponding elements in ``x``. The returned array **must** have the same shape and data type as ``x``. + 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. - - An element ``x_i`` **must** be replaced if ``x_i`` is equal to ``old`` (see :func:`~array_api.equal`) or if both ``x_i`` and ``old`` are ``NaN`` (see :func:`~array_api.isnan`). - - The function **must** consider every element of ``x``, including elements which are not explicitly stored (e.g., implicit fill values in a sparse array representation). - - The function **must not** modify ``x``. - - If ``old`` or ``new`` does not have a type and value compatible with the data type of ``x``, behavior is unspecified and thus implementation-defined. + 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): - if isinstance(arr, sp.DOK): - if copy == False: - throw(ArgumentError("cannot zero-copy change fill value of DOK format")) - arr = arr.asformat("coo") - if copy == True: - arr = copy(arr) - if arr.fill_value != old: - throw(densification warning) - replace(arr.val, copy=false) - return arr.to_dense - else: - replace(arr.val, copy=false) - arr.fill_value = new - return _agnostic.where(_agnostic.isnan(arr) if old != old else arr == old, new, arr) \ No newline at end of file + # `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) diff --git a/tests/main/test_manipulation.py b/tests/main/test_manipulation.py index f6e06d04..81c99a11 100644 --- a/tests/main/test_manipulation.py +++ b/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])