Skip to content

feat: add isin to the specification #959

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions spec/draft/API_specification/set_functions.rst
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ Objects in API
:toctree: generated
:template: method.rst

isin
unique_all
unique_counts
unique_inverse
Expand Down
48 changes: 46 additions & 2 deletions src/array_api_stubs/_draft/set_functions.py
Original file line number Diff line number Diff line change
@@ -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.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just to double-check, are we happy with e.g. torch not allowing complex values here:

In [16]: torch.isin(1j, torch.arange(3, dtype=torch.float64))
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[16], line 1
----> 1 torch.isin(1j, torch.arange(3, dtype=torch.float64))

RuntimeError: Unsupported input type encountered for isin(): ComplexDouble

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``.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So if isin(x1, x2, invert=True) is exactly equivalent to logical_not(isin(x1, x2)), we could drop the argument completely.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, I think that may be true... which in fact might be a bit awkward, because I am not sure if NaN logic adds up nicely or not.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As long as nans are never isin via equality comparison, then it seems to be unambiguous (if strange at a first sight)

In [23]: np.isin(np.nan, [np.nan], invert=True)
Out[23]: array(True)

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I suppose the weird thing is whether np.nan not in [3.] since np.nan != 3. so how does invert=True work? Like np.nan not in [3.] or not?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In [24]: np.isin(np.nan, [3], invert=True)
Out[24]: array(True)

In [25]: np.isin(np.nan, [3])
Out[25]: array(False)

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, sorry, mind-slip. Somehow I sometimes think just inverting can lead to weird things with NaNs, but that only works with the other comparisons not == and !=.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, here it only works because of equality comparison IIUC. Otherwise you're completely right, nans throw off logical inversion


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``.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a bit verbose, I don't think it's necessary to give explicit examples with invert here. In 100% of cases, invert=True applies logical_not to the output array.

- 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``.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same comment here about not duplicating with invert=True


- 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]:
Expand Down