diff --git a/spec/draft/API_specification/set_functions.rst b/spec/draft/API_specification/set_functions.rst index addf31e1f..cf2af9f68 100644 --- a/spec/draft/API_specification/set_functions.rst +++ b/spec/draft/API_specification/set_functions.rst @@ -18,6 +18,7 @@ Objects in API :toctree: generated :template: method.rst + isin unique_all unique_counts unique_inverse diff --git a/src/array_api_stubs/_draft/set_functions.py b/src/array_api_stubs/_draft/set_functions.py index 7fee77ecb..6bf409b72 100644 --- a/src/array_api_stubs/_draft/set_functions.py +++ b/src/array_api_stubs/_draft/set_functions.py @@ -1,7 +1,51 @@ -__all__ = ["unique_all", "unique_counts", "unique_inverse", "unique_values"] +__all__ = ["isin", "unique_all", "unique_counts", "unique_inverse", "unique_values"] -from ._types import Tuple, array +from ._types import Tuple, Union, array + + +def isin( + x1: Union[array, int, float, complex, bool], + x2: Union[array, int, float, complex, bool], + /, + *, + invert: bool = False, +) -> array: + """ + Tests whether each element in ``x1`` is in ``x2``. + + Parameters + ---------- + x1: Union[array, int, float, complex, bool] + first input array. **May** have any data type. + x2: Union[array, int, float, complex, bool] + second input array. **May** have any data type. + invert: bool + boolean indicating whether to invert the test criterion. If ``True``, the function **must** test whether each element in ``x1`` is *not* in ``x2``. If ``False``, the function **must** test whether each element in ``x1`` is in ``x2``. Default: ``False``. + + Returns + ------- + out: array + an array containing element-wise test results. The returned array **must** have the same shape as ``x1`` and **must** have a boolean data type. + + Notes + ----- + + - At least one of ``x1`` or ``x2`` **must** be an array. + + - If an element in ``x1`` is in ``x2``, the corresponding element in the output array **must** be ``True``; otherwise, the corresponding element in the output array **must** be ``False``. + + - Testing whether an element in ``x1`` corresponds to an element in ``x2`` **should** be determined based on value equality (see :func:`~array_api.equal`). For input arrays having floating-point data types, value-based equality implies the following behavior. + + - As ``nan`` values compare as ``False``, if an element in ``x1`` is ``nan`` and ``invert`` is ``False``, the corresponding element in the returned array **should** be ``False``. Otherwise, if an element in ``x1`` is ``nan`` and ``invert`` is ``True``, the corresponding element in the returned array **should** be ``True``. + - As complex floating-point values having at least one ``nan`` component compare as ``False``, if an element in ``x1`` is a complex floating-point value having one or more ``nan`` components and ``invert`` is ``False``, the corresponding element in the returned array **should** be ``False``. Otherwise, if an element in ``x1`` is a complex floating-point value having one or more ``nan`` components and ``invert`` is ``True``, the corresponding element in the returned array **should** be ``True``. + - As ``-0`` and ``+0`` compare as ``True``, if an element in ``x1`` is ``±0`` and ``x2`` contains at least one element which is ``±0`` + + - if ``invert`` is ``False``, the corresponding element in the returned array **should** be ``True``. + - if ``invert`` is ``True``, the corresponding element in the returned array **should** be ``False``. + + - Comparison of arrays without a corresponding promotable data type (see :ref:`type-promotion`) is unspecified and thus implementation-defined. + """ def unique_all(x: array, /) -> Tuple[array, array, array, array]: