array-api-extra 0.2.0__py3-none-any.whl → 0.2.1.dev0__py3-none-any.whl

array_api_extra/__init__.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
-from ._funcs import atleast_nd, cov, create_diagonal, expand_dims, kron, sinc
+from ._funcs import atleast_nd, cov, create_diagonal, expand_dims, kron, setdiff1d, sinc
 
-__version__ = "0.2.0"
+__version__ = "0.2.1.dev0"
 
+# pylint: disable=duplicate-code
 __all__ = [
     "__version__",
     "atleast_nd",
@@ -11,5 +12,6 @@ __all__ = [
     "create_diagonal",
     "expand_dims",
     "kron",
+    "setdiff1d",
     "sinc",
 ]

array_api_extra/_funcs.py

@@ -1,12 +1,22 @@
-from __future__ import annotations
+from __future__ import annotations  # https://github.com/pylint-dev/pylint/pull/9990
 
+import typing
 import warnings
-from typing import TYPE_CHECKING
 
-if TYPE_CHECKING:
-    from ._typing import Array, ModuleType
+if typing.TYPE_CHECKING:
+    from ._lib._typing import Array, ModuleType
 
-__all__ = ["atleast_nd", "cov", "create_diagonal", "expand_dims", "kron", "sinc"]
+from ._lib import _utils
+
+__all__ = [
+    "atleast_nd",
+    "cov",
+    "create_diagonal",
+    "expand_dims",
+    "kron",
+    "setdiff1d",
+    "sinc",
+]
 
 
 def atleast_nd(x: Array, /, *, ndim: int, xp: ModuleType) -> Array:
@@ -399,6 +409,53 @@ def kron(a: Array, b: Array, /, *, xp: ModuleType) -> Array:
     return xp.reshape(result, tuple(xp.multiply(a_shape, b_shape)))
 
 
+def setdiff1d(
+    x1: Array, x2: Array, /, *, assume_unique: bool = False, xp: ModuleType
+) -> Array:
+    """
+    Find the set difference of two arrays.
+
+    Return the unique values in `x1` that are not in `x2`.
+
+    Parameters
+    ----------
+    x1 : array
+        Input array.
+    x2 : array
+        Input comparison array.
+    assume_unique : bool
+        If ``True``, the input arrays are both assumed to be unique, which
+        can speed up the calculation. Default is ``False``.
+    xp : array_namespace
+        The standard-compatible namespace for `x1` and `x2`.
+
+    Returns
+    -------
+    res : array
+        1D array of values in `x1` that are not in `x2`. The result
+        is sorted when `assume_unique` is ``False``, but otherwise only sorted
+        if the input is sorted.
+
+    Examples
+    --------
+    >>> import array_api_strict as xp
+    >>> import array_api_extra as xpx
+
+    >>> x1 = xp.asarray([1, 2, 3, 2, 4, 1])
+    >>> x2 = xp.asarray([3, 4, 5, 6])
+    >>> xpx.setdiff1d(x1, x2, xp=xp)
+    Array([1, 2], dtype=array_api_strict.int64)
+
+    """
+
+    if assume_unique:
+        x1 = xp.reshape(x1, (-1,))
+    else:
+        x1 = xp.unique_values(x1)
+        x2 = xp.unique_values(x2)
+    return x1[_utils.in1d(x1, x2, assume_unique=True, invert=True, xp=xp)]
+
+
 def sinc(x: Array, /, *, xp: ModuleType) -> Array:
     r"""
     Return the normalized sinc function.
@@ -440,10 +497,10 @@ def sinc(x: Array, /, *, xp: ModuleType) -> Array:
 
     References
     ----------
-    .. [1] Weisstein, Eric W. "Sinc Function." From MathWorld--A Wolfram Web
-           Resource. https://mathworld.wolfram.com/SincFunction.html
-    .. [2] Wikipedia, "Sinc function",
-           https://en.wikipedia.org/wiki/Sinc_function
+    #. Weisstein, Eric W. "Sinc Function." From MathWorld--A Wolfram Web
+       Resource. https://mathworld.wolfram.com/SincFunction.html
+    #. Wikipedia, "Sinc function",
+       https://en.wikipedia.org/wiki/Sinc_function
 
     Examples
     --------

array_api_extra/_lib/_compat.py

@@ -0,0 +1,168 @@
+### Helpers borrowed from array-api-compat
+
+from __future__ import annotations  # https://github.com/pylint-dev/pylint/pull/9990
+
+import inspect
+import sys
+import typing
+
+from typing_extensions import override
+
+if typing.TYPE_CHECKING:
+    from ._typing import Array, Device
+
+__all__ = ["device"]
+
+
+# Placeholder object to represent the dask device
+# when the array backend is not the CPU.
+# (since it is not easy to tell which device a dask array is on)
+class _dask_device:  # pylint: disable=invalid-name
+    @override
+    def __repr__(self) -> str:
+        return "DASK_DEVICE"
+
+
+_DASK_DEVICE = _dask_device()
+
+
+# device() is not on numpy.ndarray or dask.array and to_device() is not on numpy.ndarray
+# or cupy.ndarray. They are not included in array objects of this library
+# because this library just reuses the respective ndarray classes without
+# wrapping or subclassing them. These helper functions can be used instead of
+# the wrapper functions for libraries that need to support both NumPy/CuPy and
+# other libraries that use devices.
+def device(x: Array, /) -> Device:
+    """
+    Hardware device the array data resides on.
+
+    This is equivalent to `x.device` according to the `standard
+    <https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.device.html>`__.
+    This helper is included because some array libraries either do not have
+    the `device` attribute or include it with an incompatible API.
+
+    Parameters
+    ----------
+    x: array
+        array instance from an array API compatible library.
+
+    Returns
+    -------
+    out: device
+        a ``device`` object (see the `Device Support <https://data-apis.org/array-api/latest/design_topics/device_support.html>`__
+        section of the array API specification).
+
+    Notes
+    -----
+
+    For NumPy the device is always `"cpu"`. For Dask, the device is always a
+    special `DASK_DEVICE` object.
+
+    See Also
+    --------
+
+    to_device : Move array data to a different device.
+
+    """
+    if _is_numpy_array(x):
+        return "cpu"
+    if _is_dask_array(x):
+        # Peek at the metadata of the jax array to determine type
+        try:
+            import numpy as np  # pylint: disable=import-outside-toplevel
+
+            if isinstance(x._meta, np.ndarray):  # pylint: disable=protected-access
+                # Must be on CPU since backed by numpy
+                return "cpu"
+        except ImportError:
+            pass
+        return _DASK_DEVICE
+    if _is_jax_array(x):
+        # JAX has .device() as a method, but it is being deprecated so that it
+        # can become a property, in accordance with the standard. In order for
+        # this function to not break when JAX makes the flip, we check for
+        # both here.
+        if inspect.ismethod(x.device):
+            return x.device()
+        return x.device
+    if _is_pydata_sparse_array(x):
+        # `sparse` will gain `.device`, so check for this first.
+        x_device = getattr(x, "device", None)
+        if x_device is not None:
+            return x_device
+        # Everything but DOK has this attr.
+        try:
+            inner = x.data
+        except AttributeError:
+            return "cpu"
+        # Return the device of the constituent array
+        return device(inner)
+    return x.device
+
+
+def _is_numpy_array(x: Array) -> bool:
+    """Return True if `x` is a NumPy array."""
+    # Avoid importing NumPy if it isn't already
+    if "numpy" not in sys.modules:
+        return False
+
+    import numpy as np  # pylint: disable=import-outside-toplevel
+
+    # TODO: Should we reject ndarray subclasses?
+    return isinstance(x, (np.ndarray, np.generic)) and not _is_jax_zero_gradient_array(
+        x
+    )
+
+
+def _is_dask_array(x: Array) -> bool:
+    """Return True if `x` is a dask.array Array."""
+    # Avoid importing dask if it isn't already
+    if "dask.array" not in sys.modules:
+        return False
+
+    # pylint: disable=import-error, import-outside-toplevel
+    import dask.array  # type: ignore[import-not-found]  # pyright: ignore[reportMissingImports]
+
+    return isinstance(x, dask.array.Array)
+
+
+def _is_jax_zero_gradient_array(x: Array) -> bool:
+    """Return True if `x` is a zero-gradient array.
+
+    These arrays are a design quirk of Jax that may one day be removed.
+    See https://github.com/google/jax/issues/20620.
+    """
+    if "numpy" not in sys.modules or "jax" not in sys.modules:
+        return False
+
+    # pylint: disable=import-error, import-outside-toplevel
+    import jax  # type: ignore[import-not-found]  # pyright: ignore[reportMissingImports]
+    import numpy as np  # pylint: disable=import-outside-toplevel
+
+    return isinstance(x, np.ndarray) and x.dtype == jax.float0  # pyright: ignore[reportUnknownVariableType]
+
+
+def _is_jax_array(x: Array) -> bool:
+    """Return True if `x` is a JAX array."""
+    # Avoid importing jax if it isn't already
+    if "jax" not in sys.modules:
+        return False
+
+    # pylint: disable=import-error, import-outside-toplevel
+    import jax  # pyright: ignore[reportMissingImports]
+
+    return isinstance(x, jax.Array) or _is_jax_zero_gradient_array(x)
+
+
+def _is_pydata_sparse_array(x: Array) -> bool:
+    """Return True if `x` is an array from the `sparse` package."""
+
+    # Avoid importing jax if it isn't already
+    if "sparse" not in sys.modules:
+        return False
+
+    # pylint: disable=import-error, import-outside-toplevel
+    import sparse  # type: ignore[import-not-found]  # pyright: ignore[reportMissingImports]
+
+    # TODO: Account for other backends.
+    return isinstance(x, sparse.SparseArray)

array_api_extra/_lib/_typing.py

@@ -0,0 +1,10 @@
+from __future__ import annotations  # https://github.com/pylint-dev/pylint/pull/9990
+
+from types import ModuleType
+from typing import Any
+
+# To be changed to a Protocol later (see data-apis/array-api#589)
+Array = Any  # type: ignore[no-any-explicit]
+Device = Any  # type: ignore[no-any-explicit]
+
+__all__ = ["Array", "Device", "ModuleType"]

array_api_extra/_lib/_utils.py

@@ -0,0 +1,65 @@
+from __future__ import annotations  # https://github.com/pylint-dev/pylint/pull/9990
+
+import typing
+
+if typing.TYPE_CHECKING:
+    from ._typing import Array, ModuleType
+
+from . import _compat
+
+__all__ = ["in1d"]
+
+
+def in1d(
+    x1: Array,
+    x2: Array,
+    /,
+    *,
+    assume_unique: bool = False,
+    invert: bool = False,
+    xp: ModuleType,
+) -> Array:
+    """Checks whether each element of an array is also present in a
+    second array.
+
+    Returns a boolean array the same length as `x1` that is True
+    where an element of `x1` is in `x2` and False otherwise.
+
+    This function has been adapted using the original implementation
+    present in numpy:
+    https://github.com/numpy/numpy/blob/v1.26.0/numpy/lib/arraysetops.py#L524-L758
+    """
+
+    # This code is run to make the code significantly faster
+    if x2.shape[0] < 10 * x1.shape[0] ** 0.145:
+        if invert:
+            mask = xp.ones(x1.shape[0], dtype=xp.bool, device=x1.device)
+            for a in x2:
+                mask &= x1 != a
+        else:
+            mask = xp.zeros(x1.shape[0], dtype=xp.bool, device=x1.device)
+            for a in x2:
+                mask |= x1 == a
+        return mask
+
+    rev_idx = xp.empty(0)  # placeholder
+    if not assume_unique:
+        x1, rev_idx = xp.unique_inverse(x1)
+        x2 = xp.unique_values(x2)
+
+    ar = xp.concat((x1, x2))
+    device_ = _compat.device(ar)
+    # We need this to be a stable sort.
+    order = xp.argsort(ar, stable=True)
+    reverse_order = xp.argsort(order, stable=True)
+    sar = xp.take(ar, order, axis=0)
+    if sar.size >= 1:
+        bool_ar = sar[1:] != sar[:-1] if invert else sar[1:] == sar[:-1]
+    else:
+        bool_ar = xp.asarray([False]) if invert else xp.asarray([True])
+    flag = xp.concat((bool_ar, xp.asarray([invert], device=device_)))
+    ret = xp.take(flag, reverse_order, axis=0)
+
+    if assume_unique:
+        return ret[: x1.shape[0]]
+    return xp.take(ret, rev_idx, axis=0)

{array_api_extra-0.2.0.dist-info → array_api_extra-0.2.1.dev0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: array-api-extra
-Version: 0.2.0
+Version: 0.2.1.dev0
 Summary: Extra array functions built on top of the array API standard.
 Project-URL: Homepage, https://github.com/data-apis/array-api-extra
 Project-URL: Bug Tracker, https://github.com/data-apis/array-api-extra/issues
@@ -41,6 +41,7 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Typing :: Typed
 Requires-Python: >=3.10
+Requires-Dist: typing-extensions
 Provides-Extra: docs
 Requires-Dist: furo>=2023.08.17; extra == 'docs'
 Requires-Dist: myst-parser>=0.13; extra == 'docs'
@@ -60,6 +61,8 @@ Description-Content-Type: text/markdown
 [![docs - here!][docs-badge]][docs-link]
 [![Pixi Badge](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/prefix-dev/pixi/main/assets/badge/v0.json)](https://pixi.sh)
 [![All Contributors](https://img.shields.io/github/all-contributors/data-apis/array-api-extra?color=ee8449&style=flat-square)](#contributors)
+[![Scientific Python Ecosystem Coordination](https://img.shields.io/badge/SPEC-0-green?labelColor=%23004811&color=%235CA038)](https://scientific-python.org/specs/)
+[![codecov](https://codecov.io/github/data-apis/array-api-extra/graph/badge.svg)](https://codecov.io/github/data-apis/array-api-extra)
 
 [![PyPI version][pypi-version]][pypi-link]
 [![Conda-Forge][conda-badge]][conda-link]
@@ -84,15 +87,24 @@ Description-Content-Type: text/markdown
 
 Extra array functions built on top of the array API standard.
 
+Used by:
+
+- [SciPy](https://github.com/scipy/scipy) - Fundamental algorithms for
+  scientific computing.
+- ...
+
 ## Contributors
 
+This project exists thanks to the following contributors
+([emoji key](https://allcontributors.org/docs/en/emoji-key)):
+
 <!-- ALL-CONTRIBUTORS-LIST:START - Do not remove or modify this section -->
 <!-- prettier-ignore-start -->
 <!-- markdownlint-disable -->
 <table>
   <tbody>
     <tr>
-      <td align="center" valign="top" width="14.28%"><a href="https://lucascolley.github.io/"><img src="https://avatars.githubusercontent.com/u/51488791?v=4?s=100" width="100px;" alt="Lucas Colley"/><br /><sub><b>Lucas Colley</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/commits?author=lucascolley" title="Code">💻</a> <a href="https://github.com/data-apis/array-api-extra/commits?author=lucascolley" title="Documentation">📖</a> <a href="#example-lucascolley" title="Examples">💡</a> <a href="#ideas-lucascolley" title="Ideas, Planning, & Feedback">🤔</a> <a href="#infra-lucascolley" title="Infrastructure (Hosting, Build-Tools, etc)">🚇</a> <a href="#maintenance-lucascolley" title="Maintenance">🚧</a> <a href="#tool-lucascolley" title="Tools">🔧</a> <a href="https://github.com/data-apis/array-api-extra/commits?author=lucascolley" title="Tests">⚠️</a> <a href="https://github.com/data-apis/array-api-extra/issues?q=author%3Alucascolley" title="Bug reports">🐛</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://lucascolley.github.io/"><img src="https://avatars.githubusercontent.com/u/51488791?v=4?s=100" width="100px;" alt="Lucas Colley"/><br /><sub><b>Lucas Colley</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/commits?author=lucascolley" title="Code">💻</a> <a href="https://github.com/data-apis/array-api-extra/commits?author=lucascolley" title="Documentation">📖</a> <a href="#example-lucascolley" title="Examples">💡</a> <a href="#ideas-lucascolley" title="Ideas, Planning, & Feedback">🤔</a> <a href="#infra-lucascolley" title="Infrastructure (Hosting, Build-Tools, etc)">🚇</a> <a href="#maintenance-lucascolley" title="Maintenance">🚧</a> <a href="#tool-lucascolley" title="Tools">🔧</a> <a href="https://github.com/data-apis/array-api-extra/commits?author=lucascolley" title="Tests">⚠️</a> <a href="https://github.com/data-apis/array-api-extra/issues?q=author%3Alucascolley" title="Bug reports">🐛</a> <a href="#platform-lucascolley" title="Packaging/porting to new platform">📦</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/mdhaber"><img src="https://avatars.githubusercontent.com/u/6570539?v=4?s=100" width="100px;" alt="Matt Haberland"/><br /><sub><b>Matt Haberland</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/commits?author=mdhaber" title="Code">💻</a> <a href="#ideas-mdhaber" title="Ideas, Planning, & Feedback">🤔</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/j-bowhay"><img src="https://avatars.githubusercontent.com/u/60778417?v=4?s=100" width="100px;" alt="Jake Bowhay"/><br /><sub><b>Jake Bowhay</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/commits?author=j-bowhay" title="Code">💻</a> <a href="https://github.com/data-apis/array-api-extra/pulls?q=is%3Apr+reviewed-by%3Aj-bowhay" title="Reviewed Pull Requests">👀</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/asmeurer"><img src="https://avatars.githubusercontent.com/u/71486?v=4?s=100" width="100px;" alt="Aaron Meurer"/><br /><sub><b>Aaron Meurer</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/pulls?q=is%3Apr+reviewed-by%3Aasmeurer" title="Reviewed Pull Requests">👀</a> <a href="#ideas-asmeurer" title="Ideas, Planning, & Feedback">🤔</a> <a href="#infra-asmeurer" title="Infrastructure (Hosting, Build-Tools, etc)">🚇</a></td>
@@ -107,6 +119,10 @@ Extra array functions built on top of the array API standard.
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/izaid"><img src="https://avatars.githubusercontent.com/u/482179?v=4?s=100" width="100px;" alt="Irwin Zaid"/><br /><sub><b>Irwin Zaid</b></sub></a><br /><a href="#ideas-izaid" title="Ideas, Planning, & Feedback">🤔</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/jakirkham"><img src="https://avatars.githubusercontent.com/u/3019665?v=4?s=100" width="100px;" alt="jakirkham"/><br /><sub><b>jakirkham</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/commits?author=jakirkham" title="Code">💻</a> <a href="https://github.com/data-apis/array-api-extra/pulls?q=is%3Apr+reviewed-by%3Ajakirkham" title="Reviewed Pull Requests">👀</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/tupui"><img src="https://avatars.githubusercontent.com/u/23188539?v=4?s=100" width="100px;" alt="Pamphile Roy"/><br /><sub><b>Pamphile Roy</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/commits?author=tupui" title="Code">💻</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://github.com/OmarManzoor"><img src="https://avatars.githubusercontent.com/u/17495884?v=4?s=100" width="100px;" alt="Omar Salman"/><br /><sub><b>Omar Salman</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/commits?author=OmarManzoor" title="Code">💻</a></td>
+    </tr>
+    <tr>
+      <td align="center" valign="top" width="14.28%"><a href="https://ogrisel.com"><img src="https://avatars.githubusercontent.com/u/89061?v=4?s=100" width="100px;" alt="Olivier Grisel"/><br /><sub><b>Olivier Grisel</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/commits?author=ogrisel" title="Code">💻</a></td>
     </tr>
   </tbody>
 </table>
@@ -115,3 +131,7 @@ Extra array functions built on top of the array API standard.
 <!-- prettier-ignore-end -->
 
 <!-- ALL-CONTRIBUTORS-LIST:END -->
+
+We follow the
+[all-contributors](https://github.com/all-contributors/all-contributors)
+specification. Contributions of any kind welcome!

array_api_extra-0.2.1.dev0.dist-info/RECORD

@@ -0,0 +1,10 @@
+array_api_extra/__init__.py,sha256=JVGamml3JhJqEAzkjNa78kXLp1nabQxjvKKoSAvBtU0,332
+array_api_extra/_funcs.py,sha256=o68LFe2YhygPGoRoe375oLfCmMQiuEmPSyaKlHkZABg,16290
+array_api_extra/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+array_api_extra/_lib/_compat.py,sha256=eneqQria9CVCPUw1HrtiV0s48vCSnAxeb32dYXuYiHE,5586
+array_api_extra/_lib/_typing.py,sha256=tJjhePbSepxiKfU6ap71gM81deKAm9rfL7wMIHqjuFs,341
+array_api_extra/_lib/_utils.py,sha256=QoUXAUzLdyNIlU76CNhPvGdgob6oDCzZwb1x0fJJABI,1948
+array_api_extra-0.2.1.dev0.dist-info/METADATA,sha256=EEK5Ptp1VQP7wfpVJWO-NanCcubEw6AS9eZvTFq5XWI,12192
+array_api_extra-0.2.1.dev0.dist-info/WHEEL,sha256=C2FUgwZgiLbznR-k0b_5k3Ai_1aASOXDss3lzCUsUug,87
+array_api_extra-0.2.1.dev0.dist-info/licenses/LICENSE,sha256=WElDmP4Uf9znamiy3s1MCM46HqI3ttZ4UAHBX4IsbtY,1097
+array_api_extra-0.2.1.dev0.dist-info/RECORD,,

array_api_extra/_typing.py

@@ -1,8 +0,0 @@
-from __future__ import annotations
-
-from types import ModuleType
-from typing import Any
-
-Array = Any  # To be changed to a Protocol later (see array-api#589)
-
-__all__ = ["Array", "ModuleType"]

array_api_extra-0.2.0.dist-info/RECORD

@@ -1,8 +0,0 @@
-array_api_extra/__init__.py,sha256=916j5GLpulyZZsUQa-I_r510XDVbap_aIrVpCVn_PIk,266
-array_api_extra/_funcs.py,sha256=T5nPgBxYOb8DkNHlEM52Qf70Nf7Qb6lFtlDtuvmEk4c,14906
-array_api_extra/_typing.py,sha256=E3XJz5PbjXP-ckQMQLi_nOJPLr-B0cm_EVArRwY-7FY,193
-array_api_extra/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-array_api_extra-0.2.0.dist-info/METADATA,sha256=ng-cZog84QSYi9U3zZC5vGM9FQqr01orTi_P3Dsylj8,10691
-array_api_extra-0.2.0.dist-info/WHEEL,sha256=C2FUgwZgiLbznR-k0b_5k3Ai_1aASOXDss3lzCUsUug,87
-array_api_extra-0.2.0.dist-info/licenses/LICENSE,sha256=WElDmP4Uf9znamiy3s1MCM46HqI3ttZ4UAHBX4IsbtY,1097
-array_api_extra-0.2.0.dist-info/RECORD,,