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.

Sign up for GitHub

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 5 commits into
base: main
Choose a base branch
from
+46 −2
Open

feat: add isin to the specification #959

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
7c09df3
feat: add `isin` to the specification
kgryte Jun 12, 2025
7259ac7
docs: fix typo
kgryte Jun 12, 2025
428be60
fix: import missing type
kgryte Jun 12, 2025
11662a4
docs: update copy
kgryte Jun 23, 2025
1c78599
docs: s/should/must/
kgryte Jun 23, 2025
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
View file
Open in desktop
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
47 changes: 45 additions & 2 deletions src/array_api_stubs/_draft/set_functions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,50 @@
__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],
Copy link
Contributor

Choose a reason for hiding this comment

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

Apologies for the very late comment! I'm afraid allowing x2 to be an array is problematic for ndonnx (or rather ONNX in general). If x2 is known "at build time" - as it would be if it were, for instance, a Sequence[int | float | complex | bool] - then there are efficient ways to implement this in the ONNX standard using a HashMap. However, if the values of x2 are not known ahead of time, then one would be forced to implement this using broadcasting, such as x1[:, None] == x2[None, :] and a subsequent any reduction. The performance of such an implementation may be surprisingly bad in some circumstances. This is also the reason why we use a Sequence rather than an array in the current ndonnx implementation.

/,
*,
invert: bool = False,
) -> array:
"""
Tests whether each element in ``x1`` is in ``x2``.
Copy link
Contributor

Choose a reason for hiding this comment

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

How about this wording:

Suggested change
Tests whether each element in ``x1`` is in ``x2``.
Tests for each element in ``x1`` if it 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

Copy link
Contributor Author

Choose a reason for hiding this comment

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

@ev-br How difficult would this be to work around in the compat layer?

Copy link
Member

Choose a reason for hiding this comment

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

It's doable of course. That said, we are then adding to the growing "thickness" of the allegedly thin compat shim that is array-api-compat, and we'd be adding one more thing that realistically pytorch is not going to implement in a foreseeable future. Meaning we should be very clear that these small steps all move the compat layer from a temporary solution to permanent. Which looks like we should just stop pretending that the compat layer is temporary.

Copy link
Contributor

Choose a reason for hiding this comment

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

What are actual use-cases for isin? I suspect it's used on floats, but in general it is more of an integer API (since true floats are likely to not match exactly), so if there are very few use-cases maybe it makes sense to limit what is promised to work?

(In the sense, that one point of the Array API was to not add a lot of complicated/awkward API.)

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

Copy link
Contributor Author

Choose a reason for hiding this comment

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

@ev-br Yes, in principle, we could drop the invert kwarg; however, the point of having that kwarg is to allow for more efficient operations in libraries not supporting graph-based optimization. Personally, I would prefer a separate "is not in" API, but the current lay of the land is having a kwarg in isin which negates the element-wise result.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I recognize that the performance optimization argument I just stated is a bit tenuous here, given that the OP explicitly advocates against including assume_unique, but the difference between invert and assume_unique is that the former conveys semantic meaning, whereas the latter is primarily about making implementer lives easier.

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`` **must** 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. When ``invert`` is ``False``,
- As ``nan`` values compare as ``False``, if an element in ``x1`` is ``nan``, the corresponding element in the returned array **must** be ``False``.
- 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, the corresponding element in the returned array **must** be ``False``.
- As ``-0`` and ``+0`` compare as ``True``, if an element in ``x1`` is ``±0`` and ``x2`` contains at least one element which is ``±0``, the corresponding element in the returned array **must** be ``True``.
When ``invert`` is ``True``, the returned array **must** contain the same results as if the operation is implemented as ``logical_not(isin(x1, x2))``.
- 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