diff --git a/spec/draft/API_specification/utility_functions.rst b/spec/draft/API_specification/utility_functions.rst index a09c99f79..e0a6f6942 100644 --- a/spec/draft/API_specification/utility_functions.rst +++ b/spec/draft/API_specification/utility_functions.rst @@ -21,3 +21,4 @@ Objects in API all any diff + replace diff --git a/src/array_api_stubs/_draft/utility_functions.py b/src/array_api_stubs/_draft/utility_functions.py index 4ac290117..5b172dc37 100644 --- a/src/array_api_stubs/_draft/utility_functions.py +++ b/src/array_api_stubs/_draft/utility_functions.py @@ -1,4 +1,4 @@ -__all__ = ["all", "any", "diff"] +__all__ = ["all", "any", "diff", "replace"] from ._types import Optional, Tuple, Union, array @@ -119,3 +119,36 @@ def diff( .. versionadded:: 2024.12 """ + + +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`). + + 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. + """