_helpers.py

"""
Various helper functions which are not part of the spec.

Functions which start with an underscore are for internal use only but helpers
that are in __all__ are intended as additional helper functions for use by end
users of the compat library.
"""

from __future__ import annotations

import enum
import inspect
import math
import sys
import warnings
from collections.abc import Collection, Hashable
from functools import lru_cache
from typing import (
    TYPE_CHECKING,
    Any,
    Final,
    Literal,
    SupportsIndex,
    TypeAlias,
    TypeGuard,
    cast,
    overload,
)

from ._typing import Array, Device, HasShape, Namespace, SupportsArrayNamespace

if TYPE_CHECKING:
    import cupy as cp
    import dask.array as da
    import jax
    import ndonnx as ndx
    import numpy as np
    import numpy.typing as npt
    import sparse
    import torch

    # TODO: import from typing (requires Python >=3.13)
    from typing_extensions import TypeIs

    _ZeroGradientArray: TypeAlias = npt.NDArray[np.void]

    _ArrayApiObj: TypeAlias = (
        npt.NDArray[Any]
        | cp.ndarray
        | da.Array
        | jax.Array
        | ndx.Array
        | sparse.SparseArray
        | torch.Tensor
        | SupportsArrayNamespace[Any]
    )

_API_VERSIONS_OLD: Final = frozenset({"2021.12", "2022.12", "2023.12"})
_API_VERSIONS: Final = _API_VERSIONS_OLD | frozenset({"2024.12"})


@lru_cache(100)
def _issubclass_fast(cls: type, modname: str, clsname: str) -> bool:
    try:
        mod = sys.modules[modname]
    except KeyError:
        return False
    parent_cls = getattr(mod, clsname)
    return issubclass(cls, parent_cls)


def _is_jax_zero_gradient_array(x: object) -> TypeGuard[_ZeroGradientArray]:
    """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.
    """
    # Fast exit
    try:
        dtype = x.dtype  # type: ignore[attr-defined]
    except AttributeError:
        return False
    cls = cast(Hashable, type(dtype))
    if not _issubclass_fast(cls, "numpy.dtypes", "VoidDType"):
        return False

    if "jax" not in sys.modules:
        return False

    import jax
    # jax.float0 is a np.dtype([('float0', 'V')])
    return dtype == jax.float0


def is_numpy_array(x: object) -> TypeIs[npt.NDArray[Any]]:
    """
    Return True if `x` is a NumPy array.

    This function does not import NumPy if it has not already been imported
    and is therefore cheap to use.

    This also returns True for `ndarray` subclasses and NumPy scalar objects.

    See Also
    --------

    array_namespace
    is_array_api_obj
    is_cupy_array
    is_torch_array
    is_ndonnx_array
    is_dask_array
    is_jax_array
    is_pydata_sparse_array
    """
    # TODO: Should we reject ndarray subclasses?
    cls = cast(Hashable, type(x))
    return (
        _issubclass_fast(cls, "numpy", "ndarray") 
        or _issubclass_fast(cls, "numpy", "generic")
    ) and not _is_jax_zero_gradient_array(x)


def is_cupy_array(x: object) -> bool:
    """
    Return True if `x` is a CuPy array.

    This function does not import CuPy if it has not already been imported
    and is therefore cheap to use.

    This also returns True for `cupy.ndarray` subclasses and CuPy scalar objects.

    See Also
    --------

    array_namespace
    is_array_api_obj
    is_numpy_array
    is_torch_array
    is_ndonnx_array
    is_dask_array
    is_jax_array
    is_pydata_sparse_array
    """
    cls = cast(Hashable, type(x))
    return _issubclass_fast(cls, "cupy", "ndarray")


def is_torch_array(x: object) -> TypeIs[torch.Tensor]:
    """
    Return True if `x` is a PyTorch tensor.

    This function does not import PyTorch if it has not already been imported
    and is therefore cheap to use.

    See Also
    --------

    array_namespace
    is_array_api_obj
    is_numpy_array
    is_cupy_array
    is_dask_array
    is_jax_array
    is_pydata_sparse_array
    """
    cls = cast(Hashable, type(x))
    return _issubclass_fast(cls, "torch", "Tensor")


def is_ndonnx_array(x: object) -> TypeIs[ndx.Array]:
    """
    Return True if `x` is a ndonnx Array.

    This function does not import ndonnx if it has not already been imported
    and is therefore cheap to use.

    See Also
    --------

    array_namespace
    is_array_api_obj
    is_numpy_array
    is_cupy_array
    is_ndonnx_array
    is_dask_array
    is_jax_array
    is_pydata_sparse_array
    """
    cls = cast(Hashable, type(x))
    return _issubclass_fast(cls, "ndonnx", "Array")


def is_dask_array(x: object) -> TypeIs[da.Array]:
    """
    Return True if `x` is a dask.array Array.

    This function does not import dask if it has not already been imported
    and is therefore cheap to use.

    See Also
    --------

    array_namespace
    is_array_api_obj
    is_numpy_array
    is_cupy_array
    is_torch_array
    is_ndonnx_array
    is_jax_array
    is_pydata_sparse_array
    """
    cls = cast(Hashable, type(x))
    return _issubclass_fast(cls, "dask.array", "Array")


def is_jax_array(x: object) -> TypeIs[jax.Array]:
    """
    Return True if `x` is a JAX array.

    This function does not import JAX if it has not already been imported
    and is therefore cheap to use.


    See Also
    --------

    array_namespace
    is_array_api_obj
    is_numpy_array
    is_cupy_array
    is_torch_array
    is_ndonnx_array
    is_dask_array
    is_pydata_sparse_array
    """
    cls = cast(Hashable, type(x))
    return _issubclass_fast(cls, "jax", "Array") or _is_jax_zero_gradient_array(x)


def is_pydata_sparse_array(x: object) -> TypeIs[sparse.SparseArray]:
    """
    Return True if `x` is an array from the `sparse` package.

    This function does not import `sparse` if it has not already been imported
    and is therefore cheap to use.


    See Also
    --------

    array_namespace
    is_array_api_obj
    is_numpy_array
    is_cupy_array
    is_torch_array
    is_ndonnx_array
    is_dask_array
    is_jax_array
    """
    # TODO: Account for other backends.
    cls = cast(Hashable, type(x))
    return _issubclass_fast(cls, "sparse", "SparseArray")


def is_array_api_obj(x: object) -> TypeGuard[_ArrayApiObj]:
    """
    Return True if `x` is an array API compatible array object.

    See Also
    --------

    array_namespace
    is_numpy_array
    is_cupy_array
    is_torch_array
    is_ndonnx_array
    is_dask_array
    is_jax_array
    """
    return (
        hasattr(x, '__array_namespace__') 
        or _is_array_api_cls(cast(Hashable, type(x)))
    )


@lru_cache(100)
def _is_array_api_cls(cls: type) -> bool:
    return (
        # TODO: drop support for numpy<2 which didn't have __array_namespace__
        _issubclass_fast(cls, "numpy", "ndarray")
        or _issubclass_fast(cls, "numpy", "generic")
        or _issubclass_fast(cls, "cupy", "ndarray")
        or _issubclass_fast(cls, "torch", "Tensor")
        or _issubclass_fast(cls, "dask.array", "Array")
        or _issubclass_fast(cls, "sparse", "SparseArray")
        # TODO: drop support for jax<0.4.32 which didn't have __array_namespace__
        or _issubclass_fast(cls, "jax", "Array")
    )


def _compat_module_name() -> str:
    assert __name__.endswith(".common._helpers")
    return __name__.removesuffix(".common._helpers")


@lru_cache(100)
def is_numpy_namespace(xp: Namespace) -> bool:
    """
    Returns True if `xp` is a NumPy namespace.

    This includes both NumPy itself and the version wrapped by array-api-compat.

    See Also
    --------

    array_namespace
    is_cupy_namespace
    is_torch_namespace
    is_ndonnx_namespace
    is_dask_namespace
    is_jax_namespace
    is_pydata_sparse_namespace
    is_array_api_strict_namespace
    """
    return xp.__name__ in {"numpy", _compat_module_name() + ".numpy"}


@lru_cache(100)
def is_cupy_namespace(xp: Namespace) -> bool:
    """
    Returns True if `xp` is a CuPy namespace.

    This includes both CuPy itself and the version wrapped by array-api-compat.

    See Also
    --------

    array_namespace
    is_numpy_namespace
    is_torch_namespace
    is_ndonnx_namespace
    is_dask_namespace
    is_jax_namespace
    is_pydata_sparse_namespace
    is_array_api_strict_namespace
    """
    return xp.__name__ in {"cupy", _compat_module_name() + ".cupy"}


@lru_cache(100)
def is_torch_namespace(xp: Namespace) -> bool:
    """
    Returns True if `xp` is a PyTorch namespace.

    This includes both PyTorch itself and the version wrapped by array-api-compat.

    See Also
    --------

    array_namespace
    is_numpy_namespace
    is_cupy_namespace
    is_ndonnx_namespace
    is_dask_namespace
    is_jax_namespace
    is_pydata_sparse_namespace
    is_array_api_strict_namespace
    """
    return xp.__name__ in {"torch", _compat_module_name() + ".torch"}


def is_ndonnx_namespace(xp: Namespace) -> bool:
    """
    Returns True if `xp` is an NDONNX namespace.

    See Also
    --------

    array_namespace
    is_numpy_namespace
    is_cupy_namespace
    is_torch_namespace
    is_dask_namespace
    is_jax_namespace
    is_pydata_sparse_namespace
    is_array_api_strict_namespace
    """
    return xp.__name__ == "ndonnx"


@lru_cache(100)
def is_dask_namespace(xp: Namespace) -> bool:
    """
    Returns True if `xp` is a Dask namespace.

    This includes both ``dask.array`` itself and the version wrapped by array-api-compat.

    See Also
    --------

    array_namespace
    is_numpy_namespace
    is_cupy_namespace
    is_torch_namespace
    is_ndonnx_namespace
    is_jax_namespace
    is_pydata_sparse_namespace
    is_array_api_strict_namespace
    """
    return xp.__name__ in {"dask.array", _compat_module_name() + ".dask.array"}


def is_jax_namespace(xp: Namespace) -> bool:
    """
    Returns True if `xp` is a JAX namespace.

    This includes ``jax.numpy`` and ``jax.experimental.array_api`` which existed in
    older versions of JAX.

    See Also
    --------

    array_namespace
    is_numpy_namespace
    is_cupy_namespace
    is_torch_namespace
    is_ndonnx_namespace
    is_dask_namespace
    is_pydata_sparse_namespace
    is_array_api_strict_namespace
    """
    return xp.__name__ in {"jax.numpy", "jax.experimental.array_api"}


def is_pydata_sparse_namespace(xp: Namespace) -> bool:
    """
    Returns True if `xp` is a pydata/sparse namespace.

    See Also
    --------

    array_namespace
    is_numpy_namespace
    is_cupy_namespace
    is_torch_namespace
    is_ndonnx_namespace
    is_dask_namespace
    is_jax_namespace
    is_array_api_strict_namespace
    """
    return xp.__name__ == "sparse"


def is_array_api_strict_namespace(xp: Namespace) -> bool:
    """
    Returns True if `xp` is an array-api-strict namespace.

    See Also
    --------

    array_namespace
    is_numpy_namespace
    is_cupy_namespace
    is_torch_namespace
    is_ndonnx_namespace
    is_dask_namespace
    is_jax_namespace
    is_pydata_sparse_namespace
    """
    return xp.__name__ == "array_api_strict"


def _check_api_version(api_version: str | None) -> None:
    if api_version in _API_VERSIONS_OLD:
        warnings.warn(
            f"The {api_version} version of the array API specification was requested but the returned namespace is actually version 2024.12"
        )
    elif api_version is not None and api_version not in _API_VERSIONS:
        raise ValueError(
            "Only the 2024.12 version of the array API specification is currently supported"
        )


class _ClsToXPInfo(enum.Enum):
    SCALAR = 0
    MAYBE_JAX_ZERO_GRADIENT = 1


@lru_cache(100)
def _cls_to_namespace(
    cls: type,
    api_version: str | None,
    use_compat: bool | None,
) -> tuple[Namespace | None, _ClsToXPInfo | None]:
    if use_compat not in (None, True, False):
        raise ValueError("use_compat must be None, True, or False")
    _use_compat = use_compat in (None, True)
    cls_ = cast(Hashable, cls)  # Make mypy happy

    if (
        _issubclass_fast(cls_, "numpy", "ndarray") 
        or _issubclass_fast(cls_, "numpy", "generic")
    ):
        if use_compat is True:
            _check_api_version(api_version)
            from .. import numpy as xp
        elif use_compat is False:
            import numpy as xp  # type: ignore[no-redef]
        else:
            # NumPy 2.0+ have __array_namespace__; however they are not
            # yet fully array API compatible.
            from .. import numpy as xp  # type: ignore[no-redef]
        return xp, _ClsToXPInfo.MAYBE_JAX_ZERO_GRADIENT

    # Note: this must happen _after_ the test for np.generic,
    # because np.float64 and np.complex128 are subclasses of float and complex.
    if issubclass(cls, int | float | complex | type(None)):
        return None, _ClsToXPInfo.SCALAR

    if _issubclass_fast(cls_, "cupy", "ndarray"):
        if _use_compat:
            _check_api_version(api_version)
            from .. import cupy as xp  # type: ignore[no-redef]
        else:
            import cupy as xp  # type: ignore[no-redef]
        return xp, None

    if _issubclass_fast(cls_, "torch", "Tensor"):
        if _use_compat:
            _check_api_version(api_version)
            from .. import torch as xp  # type: ignore[no-redef]
        else:
            import torch as xp  # type: ignore[no-redef]
        return xp, None

    if _issubclass_fast(cls_, "dask.array", "Array"):
        if _use_compat:
            _check_api_version(api_version)
            from ..dask import array as xp  # type: ignore[no-redef]
        else:
            import dask.array as xp  # type: ignore[no-redef]
        return xp, None

    # Backwards compatibility for jax<0.4.32
    if _issubclass_fast(cls_, "jax", "Array"):
        return _jax_namespace(api_version, use_compat), None

    return None, None


def _jax_namespace(api_version: str | None, use_compat: bool | None) -> Namespace:
    if use_compat:
        raise ValueError("JAX does not have an array-api-compat wrapper")
    import jax.numpy as jnp
    if not hasattr(jnp, "__array_namespace_info__"):
        # JAX v0.4.32 and newer implements the array API directly in jax.numpy.
        # For older JAX versions, it is available via jax.experimental.array_api.
        # jnp.Array objects gain the __array_namespace__ method.
        import jax.experimental.array_api  # noqa: F401
    # Test api_version
    return jnp.empty(0).__array_namespace__(api_version=api_version)


def array_namespace(
    *xs: Array | complex | None,
    api_version: str | None = None,
    use_compat: bool | None = None,
) -> Namespace:
    """
    Get the array API compatible namespace for the arrays `xs`.

    Parameters
    ----------
    xs: arrays
        one or more arrays. xs can also be Python scalars (bool, int, float,
        complex, or None), which are ignored.

    api_version: str
        The newest version of the spec that you need support for (currently
        the compat library wrapped APIs support v2024.12).

    use_compat: bool or None
        If None (the default), the native namespace will be returned if it is
        already array API compatible, otherwise a compat wrapper is used. If
        True, the compat library wrapped library will be returned. If False,
        the native library namespace is returned.

    Returns
    -------

    out: namespace
        The array API compatible namespace corresponding to the arrays in `xs`.

    Raises
    ------
    TypeError
        If `xs` contains arrays from different array libraries or contains a
        non-array.


    Typical usage is to pass the arguments of a function to
    `array_namespace()` at the top of a function to get the corresponding
    array API namespace:

    .. code:: python

       def your_function(x, y):
           xp = array_api_compat.array_namespace(x, y)
           # Now use xp as the array library namespace
           return xp.mean(x, axis=0) + 2*xp.std(y, axis=0)


    Wrapped array namespaces can also be imported directly. For example,
    `array_namespace(np.array(...))` will return `array_api_compat.numpy`.
    This function will also work for any array library not wrapped by
    array-api-compat if it explicitly defines `__array_namespace__
    <https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__array_namespace__.html>`__
    (the wrapped namespace is always preferred if it exists).

    See Also
    --------

    is_array_api_obj
    is_numpy_array
    is_cupy_array
    is_torch_array
    is_dask_array
    is_jax_array
    is_pydata_sparse_array

    """
    namespaces: set[Namespace] = set()
    for x in xs:
        xp, info = _cls_to_namespace(cast(Hashable, type(x)), api_version, use_compat)
        if info is _ClsToXPInfo.SCALAR:
            continue

        if (
            info is _ClsToXPInfo.MAYBE_JAX_ZERO_GRADIENT
            and _is_jax_zero_gradient_array(x)
        ):
            xp = _jax_namespace(api_version, use_compat)

        if xp is None:
            get_ns = getattr(x, "__array_namespace__", None)
            if get_ns is None:
                raise TypeError(f"{type(x).__name__} is not a supported array type")
            if use_compat:
                raise ValueError(
                    "The given array does not have an array-api-compat wrapper"
                )
            xp = get_ns(api_version=api_version)

        namespaces.add(xp)

    try:
        (xp,) = namespaces
        return xp
    except ValueError:
        if not namespaces:
            raise TypeError(
                "array_namespace requires at least one non-scalar array input"
            )
        raise TypeError(f"Multiple namespaces for array inputs: {namespaces}")


# backwards compatibility alias
get_namespace = array_namespace


def _check_device(bare_xp: Namespace, device: Device) -> None:  # pyright: ignore[reportUnusedFunction]
    """
    Validate dummy device on device-less array backends.

    Notes
    -----
    This function is also invoked by CuPy, which does have multiple devices
    if there are multiple GPUs available.
    However, CuPy multi-device support is currently impossible
    without using the global device or a context manager:

    https://github.com/data-apis/array-api-compat/pull/293
    """
    if bare_xp is sys.modules.get("numpy"):
        if device not in ("cpu", None):
            raise ValueError(f"Unsupported device for NumPy: {device!r}")

    elif bare_xp is sys.modules.get("dask.array"):
        if device not in ("cpu", _DASK_DEVICE, None):
            raise ValueError(f"Unsupported device for Dask: {device!r}")


# 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:
    def __repr__(self) -> Literal["DASK_DEVICE"]:
        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: _ArrayApiObj, /) -> 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"
    elif is_dask_array(x):
        # Peek at the metadata of the Dask array to determine type
        if is_numpy_array(x._meta):
            # Must be on CPU since backed by numpy
            return "cpu"
        return _DASK_DEVICE
    elif is_jax_array(x):
        # FIXME Jitted JAX arrays do not have a device attribute
        #       https://github.com/jax-ml/jax/issues/26000
        #       Return None in this case. Note that this workaround breaks
        #       the standard and will result in new arrays being created on the
        #       default device instead of the same device as the input array(s).
        x_device = getattr(x, "device", None)
        # Older JAX releases had .device() as a method, which has been replaced
        # with a property in accordance with the standard.
        if inspect.ismethod(x_device):
            return x_device()
        else:
            return x_device
    elif 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  # pyright: ignore
        except AttributeError:
            return "cpu"
        # Return the device of the constituent array
        return device(inner)  # pyright: ignore
    return x.device  # type: ignore  # pyright: ignore


# Prevent shadowing, used below
_device = device


# Based on cupy.array_api.Array.to_device
def _cupy_to_device(
    x: cp.ndarray,
    device: Device,
    /,
    stream: int | Any | None = None,
) -> cp.ndarray:
    import cupy as cp

    if device == "cpu":
        # allowing us to use `to_device(x, "cpu")`
        # is useful for portable test swapping between
        # host and device backends
        return x.get()
    if not isinstance(device, cp.cuda.Device):
        raise TypeError(f"Unsupported device type {device!r}")

    if stream is None:
        with device:
            return cp.asarray(x)

    # stream can be an int as specified in __dlpack__, or a CuPy stream
    if isinstance(stream, int):
        stream = cp.cuda.ExternalStream(stream)
    elif not isinstance(stream, cp.cuda.Stream):
        raise TypeError(f"Unsupported stream type {stream!r}")

    with device, stream:
        return cp.asarray(x)


def _torch_to_device(
    x: torch.Tensor,
    device: torch.device | str | int,
    /,
    stream: int | Any | None = None,
) -> torch.Tensor:
    if stream is not None:
        raise NotImplementedError
    return x.to(device)


def to_device(x: Array, device: Device, /, *, stream: int | Any | None = None) -> Array:
    """
    Copy the array from the device on which it currently resides to the specified ``device``.

    This is equivalent to `x.to_device(device, stream=stream)` according to
    the `standard
    <https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.to_device.html>`__.
    This helper is included because some array libraries do not have the
    `to_device` method.

    Parameters
    ----------

    x: array
        array instance from an array API compatible library.

    device: 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).

    stream: int | Any | None
        stream object to use during copy. In addition to the types supported
        in ``array.__dlpack__``, implementations may choose to support any
        library-specific stream object with the caveat that any code using
        such an object would not be portable.

    Returns
    -------

    out: array
        an array with the same data and data type as ``x`` and located on the
        specified ``device``.

    Notes
    -----

    For NumPy, this function effectively does nothing since the only supported
    device is the CPU. For CuPy, this method supports CuPy CUDA
    :external+cupy:class:`Device <cupy.cuda.Device>` and
    :external+cupy:class:`Stream <cupy.cuda.Stream>` objects. For PyTorch,
    this is the same as :external+torch:meth:`x.to(device) <torch.Tensor.to>`
    (the ``stream`` argument is not supported in PyTorch).

    See Also
    --------

    device : Hardware device the array data resides on.

    """
    if is_numpy_array(x):
        if stream is not None:
            raise ValueError("The stream argument to to_device() is not supported")
        if device == "cpu":
            return x
        raise ValueError(f"Unsupported device {device!r}")
    elif is_cupy_array(x):
        # cupy does not yet have to_device
        return _cupy_to_device(x, device, stream=stream)
    elif is_torch_array(x):
        return _torch_to_device(x, device, stream=stream)
    elif is_dask_array(x):
        if stream is not None:
            raise ValueError("The stream argument to to_device() is not supported")
        # TODO: What if our array is on the GPU already?
        if device == "cpu":
            return x
        raise ValueError(f"Unsupported device {device!r}")
    elif is_jax_array(x):
        if not hasattr(x, "__array_namespace__"):
            # In JAX v0.4.31 and older, this import adds to_device method to x...
            import jax.experimental.array_api  # noqa: F401  # pyright: ignore

            # ... but only on eager JAX. It won't work inside jax.jit.
            if not hasattr(x, "to_device"):
                return x
        return x.to_device(device, stream=stream)
    elif is_pydata_sparse_array(x) and device == _device(x):
        # Perform trivial check to return the same array if
        # device is same instead of err-ing.
        return x
    return x.to_device(device, stream=stream)  # pyright: ignore


@overload
def size(x: HasShape[Collection[SupportsIndex]]) -> int: ...
@overload
def size(x: HasShape[Collection[SupportsIndex | None]]) -> int | None: ...
def size(x: HasShape[Collection[SupportsIndex | None]]) -> int | None:
    """
    Return the total number of elements of x.

    This is equivalent to `x.size` according to the `standard
    <https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.size.html>`__.

    This helper is included because PyTorch defines `size` in an
    :external+torch:meth:`incompatible way <torch.Tensor.size>`.
    It also fixes dask.array's behaviour which returns nan for unknown sizes, whereas
    the standard requires None.
    """
    # Lazy API compliant arrays, such as ndonnx, can contain None in their shape
    if None in x.shape:
        return None
    out = math.prod(cast("Collection[SupportsIndex]", x.shape))
    # dask.array.Array.shape can contain NaN
    return None if math.isnan(out) else out


@lru_cache(100)
def _is_writeable_cls(cls: type) -> bool | None:
    if (
        _issubclass_fast(cls, "numpy", "generic")
        or _issubclass_fast(cls, "jax", "Array")
        or _issubclass_fast(cls, "sparse", "SparseArray")
    ):
        return False
    if _is_array_api_cls(cls):
        return True
    return None


def is_writeable_array(x: object) -> TypeGuard[_ArrayApiObj]:
    """
    Return False if ``x.__setitem__`` is expected to raise; True otherwise.
    Return False if `x` is not an array API compatible object.

    Warning
    -------
    As there is no standard way to check if an array is writeable without actually
    writing to it, this function blindly returns True for all unknown array types.
    """
    cls = cast(Hashable, type(x))
    if _issubclass_fast(cls, "numpy", "ndarray"):
        return cast("npt.NDArray", x).flags.writeable
    res = _is_writeable_cls(cls)
    if res is not None:
        return res
    return hasattr(x, '__array_namespace__')


@lru_cache(100)
def _is_lazy_cls(cls: type) -> bool | None:
    if (
        _issubclass_fast(cls, "numpy", "ndarray")
        or _issubclass_fast(cls, "numpy", "generic")
        or _issubclass_fast(cls, "cupy", "ndarray")
        or _issubclass_fast(cls, "torch", "Tensor")
        or _issubclass_fast(cls, "sparse", "SparseArray")
    ):
        return False
    if (
        _issubclass_fast(cls, "jax", "Array")
        or _issubclass_fast(cls, "dask.array", "Array")
        or _issubclass_fast(cls, "ndonnx", "Array")
    ):
        return True
    return  None


def is_lazy_array(x: object) -> TypeGuard[_ArrayApiObj]:
    """Return True if x is potentially a future or it may be otherwise impossible or
    expensive to eagerly read its contents, regardless of their size, e.g. by
    calling ``bool(x)`` or ``float(x)``.

    Return False otherwise; e.g. ``bool(x)`` etc. is guaranteed to succeed and to be
    cheap as long as the array has the right dtype and size.

    Note
    ----
    This function errs on the side of caution for array types that may or may not be
    lazy, e.g. JAX arrays, by always returning True for them.
    """
    # **JAX note:** while it is possible to determine if you're inside or outside
    # jax.jit by testing the subclass of a jax.Array object, as well as testing bool()
    # as we do below for unknown arrays, this is not recommended by JAX best practices.

    # **Dask note:** Dask eagerly computes the graph on __bool__, __float__, and so on.
    # This behaviour, while impossible to change without breaking backwards
    # compatibility, is highly detrimental to performance as the whole graph will end
    # up being computed multiple times.

    # Note: skipping reclassification of JAX zero gradient arrays, as one will
    # exclusively get them once they leave a jax.grad JIT context.
    cls = cast(Hashable, type(x))
    res = _is_lazy_cls(cls)
    if res is not None:
        return res

    if not hasattr(x, "__array_namespace__"):
        return False

    # Unknown Array API compatible object. Note that this test may have dire consequences
    # in terms of performance, e.g. for a lazy object that eagerly computes the graph
    # on __bool__ (dask is one such example, which however is special-cased above).

    # Select a single point of the array
    s = size(cast("HasShape[Collection[SupportsIndex | None]]", x))
    if s is None:
        return True
    xp = array_namespace(x)
    if s > 1:
        x = xp.reshape(x, (-1,))[0]
    # Cast to dtype=bool and deal with size 0 arrays
    x = xp.any(x)

    try:
        bool(x)
        return False
    # The Array API standard dictactes that __bool__ should raise TypeError if the
    # output cannot be defined.
    # Here we allow for it to raise arbitrary exceptions, e.g. like Dask does.
    except Exception:
        return True


__all__ = [
    "array_namespace",
    "device",
    "get_namespace",
    "is_array_api_obj",
    "is_array_api_strict_namespace",
    "is_cupy_array",
    "is_cupy_namespace",
    "is_dask_array",
    "is_dask_namespace",
    "is_jax_array",
    "is_jax_namespace",
    "is_numpy_array",
    "is_numpy_namespace",
    "is_torch_array",
    "is_torch_namespace",
    "is_ndonnx_array",
    "is_ndonnx_namespace",
    "is_pydata_sparse_array",
    "is_pydata_sparse_namespace",
    "is_writeable_array",
    "is_lazy_array",
    "size",
    "to_device",
]

_all_ignore = ['lru_cache', 'sys', 'math', 'inspect', 'warnings']

def __dir__() -> list[str]:
    return __all__