diff --git a/meson.build b/meson.build index 042cd66d..df567ef8 100644 --- a/meson.build +++ b/meson.build @@ -10,25 +10,45 @@ py = import('python').find_installation() sources = { 'array_api_extra': files( 'src/array_api_extra/__init__.py', - 'src/array_api_extra/_delegation.py', + 'src/array_api_extra/_at.py', + 'src/array_api_extra/_creation.py', + 'src/array_api_extra/_elementwise.py', + 'src/array_api_extra/_indexing.py', + 'src/array_api_extra/_lazy.py', + 'src/array_api_extra/_linalg.py', + 'src/array_api_extra/_manipulation.py', + 'src/array_api_extra/_searching.py', + 'src/array_api_extra/_set.py', + 'src/array_api_extra/_sorting.py', + 'src/array_api_extra/_statistical.py', 'src/array_api_extra/py.typed', - 'src/array_api_extra/testing.py', + ), + 'array_api_extra/_agnostic': files( + 'src/array_api_extra/_agnostic/__init__.py', + 'src/array_api_extra/_agnostic/_creation.py', + 'src/array_api_extra/_agnostic/_elementwise.py', + 'src/array_api_extra/_agnostic/_indexing.py', + 'src/array_api_extra/_agnostic/_inspection.py', + 'src/array_api_extra/_agnostic/_linalg.py', + 'src/array_api_extra/_agnostic/_manipulation.py', + 'src/array_api_extra/_agnostic/_searching.py', + 'src/array_api_extra/_agnostic/_set.py', + 'src/array_api_extra/_agnostic/_sorting.py', + 'src/array_api_extra/_agnostic/_statistical.py', + ), + 'array_api_extra/testing': files( + 'src/array_api_extra/testing/__init__.py', + 'src/array_api_extra/testing/_testing.py', ), 'array_api_extra/_lib': files( 'src/array_api_extra/_lib/__init__.py', - 'src/array_api_extra/_lib/_at.py', 'src/array_api_extra/_lib/_backends.py', - 'src/array_api_extra/_lib/_funcs.py', - 'src/array_api_extra/_lib/_lazy.py', + 'src/array_api_extra/_lib/_compat.py', + 'src/array_api_extra/_lib/_compat.pyi', + 'src/array_api_extra/_lib/_helpers.py', 'src/array_api_extra/_lib/_testing.py', - ), - 'array_api_extra/_lib/_utils': files( - 'src/array_api_extra/_lib/_utils/__init__.py', - 'src/array_api_extra/_lib/_utils/_compat.py', - 'src/array_api_extra/_lib/_utils/_compat.pyi', - 'src/array_api_extra/_lib/_utils/_helpers.py', - 'src/array_api_extra/_lib/_utils/_typing.py', - 'src/array_api_extra/_lib/_utils/_typing.pyi', + 'src/array_api_extra/_lib/_typing.py', + 'src/array_api_extra/_lib/_typing.pyi', ), } diff --git a/src/array_api_extra/__init__.py b/src/array_api_extra/__init__.py index 0910bd44..5b4b6cfc 100644 --- a/src/array_api_extra/__init__.py +++ b/src/array_api_extra/__init__.py @@ -1,46 +1,22 @@ """Extra array functions built on top of the array API standard.""" from . import testing -from ._delegation import ( - argpartition, - atleast_nd, - broadcast_shapes, - cov, - create_diagonal, - deg2rad, - diag_indices, - expand_dims, - isclose, - isin, - kron, - nan_to_num, - nanmax, - nanmin, - nansum, - nunique, - one_hot, - pad, - partition, - rad2deg, - searchsorted, - setdiff1d, - sinc, - tril_indices, - triu_indices, - union1d, - unravel_index, -) -from ._lib._at import at -from ._lib._funcs import ( - angle, - apply_where, - default_dtype, -) -from ._lib._lazy import lazy_apply +from ._agnostic._elementwise import angle, apply_where +from ._agnostic._inspection import default_dtype +from ._at import at +from ._creation import create_diagonal, one_hot +from ._elementwise import deg2rad, isclose, nan_to_num, rad2deg, sinc +from ._indexing import diag_indices, tril_indices, triu_indices, unravel_index +from ._lazy import lazy_apply +from ._linalg import kron +from ._manipulation import atleast_nd, broadcast_shapes, expand_dims, pad +from ._searching import searchsorted +from ._set import isin, nunique, setdiff1d, union1d +from ._sorting import argpartition, partition +from ._statistical import cov, nanmax, nanmin, nansum __version__ = "0.11.2.dev0" -# pylint: disable=duplicate-code __all__ = [ "__version__", "angle", diff --git a/src/array_api_extra/_agnostic/__init__.py b/src/array_api_extra/_agnostic/__init__.py new file mode 100644 index 00000000..8af17333 --- /dev/null +++ b/src/array_api_extra/_agnostic/__init__.py @@ -0,0 +1,27 @@ +"""Array-agnostic function implementations.""" + +from . import ( + _creation, + _elementwise, + _indexing, + _inspection, + _linalg, + _manipulation, + _searching, + _set, + _sorting, + _statistical, +) + +__all__ = [ + "_creation", + "_elementwise", + "_indexing", + "_inspection", + "_linalg", + "_manipulation", + "_searching", + "_set", + "_sorting", + "_statistical", +] diff --git a/src/array_api_extra/_agnostic/_creation.py b/src/array_api_extra/_agnostic/_creation.py new file mode 100644 index 00000000..ebbd361b --- /dev/null +++ b/src/array_api_extra/_agnostic/_creation.py @@ -0,0 +1,44 @@ +"""Array-agnostic implementations for creation functions.""" + +from .._at import at +from .._lib import _compat +from .._lib._helpers import eager_shape, ndindex +from .._lib._typing import Array, ArrayNamespace + +__all__ = ["create_diagonal", "one_hot"] + + +def create_diagonal( + x: Array, /, *, offset: int = 0, xp: ArrayNamespace +) -> Array: # numpydoc ignore=PR01,RT01 + """See docstring in array_api_extra._delegation.""" + x_shape = eager_shape(x) + batch_dims = x_shape[:-1] + n = x_shape[-1] + abs(offset) + diag = xp.zeros((*batch_dims, n**2), dtype=x.dtype, device=_compat.device(x)) + + target_slice = slice( + offset if offset >= 0 else abs(offset) * n, + min(n * (n - offset), diag.shape[-1]), + n + 1, + ) + for index in ndindex(*batch_dims): + diag = at(diag)[(*index, target_slice)].set(x[(*index, slice(None))]) + return xp.reshape(diag, (*batch_dims, n, n)) + + +def one_hot( + x: Array, + /, + num_classes: int, + *, + xp: ArrayNamespace, +) -> Array: # numpydoc ignore=PR01,RT01 + """See docstring in `array_api_extra._delegation.py`.""" + # TODO: Benchmark whether this is faster on the NumPy backend: + # if is_numpy_array(x): + # out = xp.zeros((x.size, num_classes), dtype=dtype) + # out[xp.arange(x.size), xp.reshape(x, (-1,))] = 1 + # return xp.reshape(out, (*x.shape, num_classes)) + range_num_classes = xp.arange(num_classes, dtype=x.dtype, device=_compat.device(x)) + return x[..., xp.newaxis] == range_num_classes diff --git a/src/array_api_extra/_agnostic/_elementwise.py b/src/array_api_extra/_agnostic/_elementwise.py new file mode 100644 index 00000000..969e6c34 --- /dev/null +++ b/src/array_api_extra/_agnostic/_elementwise.py @@ -0,0 +1,353 @@ +"""Array-agnostic implementations for element-wise functions.""" + +from collections.abc import Callable +from types import NoneType +from typing import overload + +from .._at import at +from .._lib import _compat +from .._lib._compat import array_namespace, is_dask_namespace +from .._lib._helpers import asarrays, capabilities, meta_namespace +from .._lib._typing import Array, ArrayNamespace +from ._inspection import default_dtype + +__all__ = [ + "angle", + "apply_where", + "deg2rad", + "isclose", + "nan_to_num", + "rad2deg", + "sinc", +] + + +@overload +def apply_where( # numpydoc ignore=GL08 + cond: Array, + args: Array | tuple[Array, ...], + f1: Callable[..., Array], + f2: Callable[..., Array], + /, + *, + kwargs: dict[str, Array] | None = None, + xp: ArrayNamespace | None = None, +) -> Array: ... + + +@overload +def apply_where( # numpydoc ignore=GL08 + cond: Array, + args: Array | tuple[Array, ...], + f1: Callable[..., Array], + /, + *, + fill_value: Array | complex, + kwargs: dict[str, Array] | None = None, + xp: ArrayNamespace | None = None, +) -> Array: ... + + +def apply_where( # numpydoc ignore=PR01,PR02 + cond: Array, + args: Array | tuple[Array, ...], + f1: Callable[..., Array], + f2: Callable[..., Array] | None = None, + /, + *, + fill_value: Array | complex | None = None, + kwargs: dict[str, Array] | None = None, + xp: ArrayNamespace | None = None, +) -> Array: + """ + Run one of two elementwise functions depending on a condition. + + Equivalent to ``f1(*args) if cond else fill_value`` performed elementwise + when `fill_value` is defined, otherwise to ``f1(*args) if cond else f2(*args)``. + + Parameters + ---------- + cond : array + The condition, expressed as a boolean array. + args : Array or tuple of Arrays + Argument(s) to `f1` (and `f2`). Must be broadcastable with `cond`. + f1 : callable + Elementwise function of `args`, returning a single array. + Where `cond` is True, output will be ``f1(arg0[cond], arg1[cond], ...)``. + f2 : callable, optional + Elementwise function of `args`, returning a single array. + Where `cond` is False, output will be ``f2(arg0[cond], arg1[cond], ...)``. + Mutually exclusive with `fill_value`. + fill_value : Array or scalar, optional + If provided, value with which to fill output array where `cond` is False. + It does not need to be scalar; it needs however to be broadcastable with + `cond` and `args`. + Mutually exclusive with `f2`. You must provide one or the other. + kwargs : dict of str : Array pairs + Keyword argument(s) to `f1` (and `f2`). Values must be broadcastable with + `cond`. + xp : array_namespace, optional + The standard-compatible namespace for `cond` and `args`. Default: infer. + + Returns + ------- + Array + An array with elements from the output of `f1` where `cond` is True and either + the output of `f2` or `fill_value` where `cond` is False. The returned array has + data type determined by type promotion rules between the output of `f1` and + either `fill_value` or the output of `f2`. + + Notes + ----- + ``xp.where(cond, f1(*args), f2(*args))`` requires explicitly evaluating `f1` even + when `cond` is False, and `f2` when cond is True. This function evaluates each + function only for their matching condition, if the backend allows for it. + + On Dask, `f1` and `f2` are applied to the individual chunks and should use functions + from the namespace of the chunks. + + Examples + -------- + >>> import array_api_strict as xp + >>> import array_api_extra as xpx + >>> a = xp.asarray([5, 4, 3]) + >>> b = xp.asarray([0, 2, 2]) + >>> def f(a, b): + ... return a // b + >>> xpx.apply_where(b != 0, (a, b), f, fill_value=xp.nan) + array([ nan, 2., 1.]) + """ + # Parse and normalize arguments + if (f2 is None) == (fill_value is None): + msg = "Exactly one of `fill_value` or `f2` must be given." + raise TypeError(msg) + args_ = list(args) if isinstance(args, tuple) else [args] + del args + + kwargs_ = {} if kwargs is None else kwargs + kwkeys = list(kwargs_.keys()) + args_ = [*args_, *kwargs_.values()] + del kwargs + + xp = array_namespace(cond, fill_value, *args_) if xp is None else xp + + if isinstance(fill_value, int | float | complex | NoneType): + cond, *args_ = xp.broadcast_arrays(cond, *args_) + else: + cond, fill_value, *args_ = xp.broadcast_arrays(cond, fill_value, *args_) + + if is_dask_namespace(xp): + meta_xp = meta_namespace(cond, fill_value, *args_, xp=xp) + # map_blocks doesn't descend into tuples of Arrays + return xp.map_blocks( + _apply_where, cond, f1, f2, fill_value, *args_, kwkeys=kwkeys, xp=meta_xp + ) + + return _apply_where(cond, f1, f2, fill_value, *args_, kwkeys=kwkeys, xp=xp) + + +def _apply_where( # numpydoc ignore=PR01,RT01 + cond: Array, + f1: Callable[..., Array], + f2: Callable[..., Array] | None, + fill_value: Array | complex | bool | None, + *args: Array, + kwkeys: list[str], + xp: ArrayNamespace, +) -> Array: + """Helper of `apply_where`. On Dask, this runs on a single chunk.""" + + nargs = len(args) - len(kwkeys) + kwargs = dict(zip(kwkeys, args[nargs:], strict=True)) + args = args[:nargs] + + if not capabilities(xp, device=_compat.device(cond))["boolean indexing"]: + # jax.jit does not support assignment by boolean mask + return xp.where( + cond, + f1(*args, **kwargs), + f2(*args, **kwargs) if f2 is not None else fill_value, + ) + + temp1 = f1( + *(arr[cond] for arr in args), **{key: val[cond] for key, val in kwargs.items()} + ) + + if f2 is None: + dtype = xp.result_type(temp1, fill_value) + if isinstance(fill_value, int | float | complex): + out = xp.full_like(cond, dtype=dtype, fill_value=fill_value) + else: + out = xp.astype(fill_value, dtype, copy=True) + else: + ncond = ~cond + temp2 = f2( + *(arr[ncond] for arr in args), + **{key: val[ncond] for key, val in kwargs.items()}, + ) + dtype = xp.result_type(temp1, temp2) + out = xp.empty_like(cond, dtype=dtype) + out = at(out, ncond).set(temp2) + + return at(out, cond).set(temp1) + + +def isclose( + a: Array | complex, + b: Array | complex, + *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, + xp: ArrayNamespace, +) -> Array: # numpydoc ignore=PR01,RT01 + """See docstring in array_api_extra._delegation.""" + a, b = asarrays(a, b, xp=xp) + + a_inexact = xp.isdtype(a.dtype, ("real floating", "complex floating")) + b_inexact = xp.isdtype(b.dtype, ("real floating", "complex floating")) + if a_inexact or b_inexact: + # prevent warnings on NumPy and Dask on inf - inf + mxp = meta_namespace(a, b, xp=xp) + out = apply_where( + xp.isinf(a) | xp.isinf(b), + (a, b), + lambda a, b: mxp.isinf(a) & mxp.isinf(b) & (mxp.sign(a) == mxp.sign(b)), # pyright: ignore[reportUnknownArgumentType] + # Note: inf <= inf is True! + lambda a, b: mxp.abs(a - b) <= (atol + rtol * mxp.abs(b)), # pyright: ignore[reportUnknownArgumentType] + xp=xp, + ) + if equal_nan: + out = xp.where(xp.isnan(a) & xp.isnan(b), True, out) + return out + + if xp.isdtype(a.dtype, "bool") or xp.isdtype(b.dtype, "bool"): + if atol >= 1 or rtol >= 1: + return xp.ones_like(a == b) + return a == b + + # integer types + atol = int(atol) + if rtol == 0: + return xp.abs(a - b) <= atol + + # Don't rely on OverflowError, as it is not guaranteed by the Array API. + nrtol = int(1.0 / rtol) + if nrtol > xp.iinfo(b.dtype).max: + # rtol * max_int < 1, so it's inconsequential + return xp.abs(a - b) <= atol + return xp.abs(a - b) <= (atol + xp.abs(b) // nrtol) + + +def nan_to_num( # numpydoc ignore=PR01,RT01 + x: Array, + /, + fill_value: float = 0.0, + *, + xp: ArrayNamespace, +) -> Array: + """See docstring in `array_api_extra._delegation.py`.""" + + def perform_replacements( # numpydoc ignore=PR01,RT01 + x: Array, + fill_value: float, + xp: ArrayNamespace, + ) -> Array: + """Internal function to perform the replacements.""" + x = xp.where(xp.isnan(x), fill_value, x) + + # convert infinities to finite values + finfo = xp.finfo(x.dtype) + idx_posinf = xp.isinf(x) & ~xp.signbit(x) + idx_neginf = xp.isinf(x) & xp.signbit(x) + x = xp.where(idx_posinf, finfo.max, x) + return xp.where(idx_neginf, finfo.min, x) + + if xp.isdtype(x.dtype, "complex floating"): + return perform_replacements( + xp.real(x), + fill_value, + xp, + ) + 1j * perform_replacements( + xp.imag(x), + fill_value, + xp, + ) + + if xp.isdtype(x.dtype, "numeric"): + return perform_replacements(x, fill_value, xp) + + return x + + +def sinc(x: Array, /, *, xp: ArrayNamespace) -> Array: + # numpydoc ignore=PR01,RT01 + """See docstring in `array_api_extra._delegation.py`.""" + + # no scalars in `where` - array-api#807 + y = xp.pi * xp.where( + xp.astype(x, xp.bool), + x, + xp.asarray(xp.finfo(x.dtype).eps, dtype=x.dtype, device=_compat.device(x)), + ) + return xp.sin(y) / y + + +def angle(z: Array, /, *, deg: bool = False, xp: ArrayNamespace | None = None) -> Array: + """ + Return the angle of the complex argument. + + Parameters + ---------- + z : Array + Input array. + deg : bool, optional + Return angle in degrees if True, radians if False (default). + xp : array_namespace, optional + The standard-compatible namespace for `z`. Default: infer. + + Returns + ------- + array + The counterclockwise angle from the positive real axis on the complex + plane in the range ``(-pi, pi]``. + + Notes + ----- + Real input ``x`` is interpreted as ``x + 0j``. + + Examples + -------- + >>> import array_api_strict as xp + >>> import array_api_extra as xpx + >>> xpx.angle(xp.asarray([1.0, 1.0j, 1 + 1j]), xp=xp) + Array([0. , 1.57079633, 0.78539816], dtype=array_api_strict.float64) + >>> xpx.angle(xp.asarray([1.0, 1.0j, 1 + 1j]), deg=True, xp=xp) + Array([ 0., 90., 45.], dtype=array_api_strict.float64) + """ + if xp is None: + xp = array_namespace(z) + if xp.isdtype(z.dtype, "complex floating"): + zimag = xp.imag(z) + zreal = xp.real(z) + else: + if not xp.isdtype(z.dtype, "real floating"): + z = xp.astype(z, default_dtype(xp, device=_compat.device(z))) + zimag = xp.zeros_like(z) + zreal = z + a = xp.atan2(zimag, zreal) + if deg: + a = a * 180 / xp.pi + return a + + +def deg2rad(x: Array, /, *, xp: ArrayNamespace) -> Array: + # numpydoc ignore=PR01,RT01 + """See docstring in `array_api_extra._delegation.py`.""" + return x * xp.pi / 180 + + +def rad2deg(x: Array, /, *, xp: ArrayNamespace) -> Array: + # numpydoc ignore=PR01,RT01 + """See docstring in `array_api_extra._delegation.py`.""" + return x * 180 / xp.pi diff --git a/src/array_api_extra/_agnostic/_indexing.py b/src/array_api_extra/_agnostic/_indexing.py new file mode 100644 index 00000000..efc30942 --- /dev/null +++ b/src/array_api_extra/_agnostic/_indexing.py @@ -0,0 +1,68 @@ +"""Array-agnostic implementations for indexing functions.""" + +from .._lib._typing import Array, ArrayNamespace, Device + +__all__ = ["diag_indices", "tril_indices", "triu_indices", "unravel_index"] + + +def diag_indices( + n: int, /, *, ndim: int, device: Device | None, xp: ArrayNamespace +) -> tuple[Array, ...]: # numpydoc ignore=PR01,RT01 + """See docstring in array_api_extra._delegation.""" + idx = xp.arange(n, device=device) + return (idx,) * ndim + + +def _tri_indices( + n: int, + *, + offset: int, + m: int | None, + upper: bool, + device: Device | None, + xp: ArrayNamespace, +) -> tuple[Array, Array]: # numpydoc ignore=PR01,RT01 + """Shared implementation for `tril_indices` and `triu_indices`.""" + cols = n if m is None else m + rows = xp.arange(n, device=device)[:, xp.newaxis] + cols_a = xp.arange(cols, device=device)[xp.newaxis, :] + delta = cols_a - rows + mask = delta >= offset if upper else delta <= offset + r, c = xp.nonzero(mask) + return (r, c) + + +def tril_indices( + n: int, + /, + *, + offset: int, + m: int | None, + device: Device | None, + xp: ArrayNamespace, +) -> tuple[Array, Array]: # numpydoc ignore=PR01,RT01 + """See docstring in array_api_extra._delegation.""" + return _tri_indices(n, offset=offset, m=m, upper=False, device=device, xp=xp) + + +def triu_indices( + n: int, + /, + *, + offset: int, + m: int | None, + device: Device | None, + xp: ArrayNamespace, +) -> tuple[Array, Array]: # numpydoc ignore=PR01,RT01 + """See docstring in array_api_extra._delegation.""" + return _tri_indices(n, offset=offset, m=m, upper=True, device=device, xp=xp) + + +def unravel_index(indices: Array, shape: tuple[int, ...], /) -> tuple[Array, ...]: + # numpydoc ignore=PR01,RT01 + """See docstring in `array_api_extra._delegation.py`.""" + coords: list[Array] = [] + for dim in reversed(shape): + coords.append(indices % dim) + indices = indices // dim + return tuple(reversed(coords)) diff --git a/src/array_api_extra/_agnostic/_inspection.py b/src/array_api_extra/_agnostic/_inspection.py new file mode 100644 index 00000000..ba707d10 --- /dev/null +++ b/src/array_api_extra/_agnostic/_inspection.py @@ -0,0 +1,45 @@ +"""Array-agnostic implementations for inspection functions.""" + +from typing import Literal + +from .._lib._typing import ArrayNamespace, Device, DType + +__all__ = ["default_dtype"] + + +def default_dtype( + xp: ArrayNamespace, + kind: Literal[ + "real floating", "complex floating", "integral", "indexing" + ] = "real floating", + *, + device: Device | None = None, +) -> DType: + """ + Return the default dtype for the given namespace and device. + + This is a convenience shorthand for + ``xp.__array_namespace_info__().default_dtypes(device=device)[kind]``. + + Parameters + ---------- + xp : array_namespace + The standard-compatible namespace for which to get the default dtype. + kind : {'real floating', 'complex floating', 'integral', 'indexing'}, optional + The kind of dtype to return. Default is 'real floating'. + device : Device, optional + The device for which to get the default dtype. Default: current device. + + Returns + ------- + dtype + The default dtype for the given namespace, kind, and device. + """ + dtypes = xp.__array_namespace_info__().default_dtypes(device=device) + try: + return dtypes[kind] + except KeyError as e: + domain = ("real floating", "complex floating", "integral", "indexing") + assert set(dtypes) == set(domain), f"Non-compliant namespace: {dtypes}" + msg = f"Unknown kind '{kind}'. Expected one of {domain}." + raise ValueError(msg) from e diff --git a/src/array_api_extra/_agnostic/_linalg.py b/src/array_api_extra/_agnostic/_linalg.py new file mode 100644 index 00000000..e5f5a023 --- /dev/null +++ b/src/array_api_extra/_agnostic/_linalg.py @@ -0,0 +1,47 @@ +"""Array-agnostic implementations for linear algebra functions.""" + +from typing import cast + +from .._lib._helpers import eager_shape +from .._lib._typing import Array, ArrayNamespace +from ._manipulation import expand_dims + +__all__ = ["kron"] + + +def kron( + a: Array, + b: Array, + /, + *, + xp: ArrayNamespace, +) -> Array: # numpydoc ignore=PR01,RT01 + """See docstring in array_api_extra._delegation.""" + + singletons = (1,) * (b.ndim - a.ndim) + a = cast(Array, xp.broadcast_to(a, singletons + a.shape)) + + nd_b, nd_a = b.ndim, a.ndim + nd_max = max(nd_b, nd_a) + if nd_a == 0 or nd_b == 0: + return xp.multiply(a, b) + + a_shape = eager_shape(a) + b_shape = eager_shape(b) + + # Equalise the shapes by prepending smaller one with 1s + a_shape = (1,) * max(0, nd_b - nd_a) + a_shape + b_shape = (1,) * max(0, nd_a - nd_b) + b_shape + + # Insert empty dimensions + a_arr = expand_dims(a, axis=tuple(range(nd_b - nd_a)), xp=xp) + b_arr = expand_dims(b, axis=tuple(range(nd_a - nd_b)), xp=xp) + + # Compute the product + a_arr = expand_dims(a_arr, axis=tuple(range(1, nd_max * 2, 2)), xp=xp) + b_arr = expand_dims(b_arr, axis=tuple(range(0, nd_max * 2, 2)), xp=xp) + result = xp.multiply(a_arr, b_arr) + + # Reshape back and return + res_shape = tuple(a_s * b_s for a_s, b_s in zip(a_shape, b_shape, strict=True)) + return xp.reshape(result, res_shape) diff --git a/src/array_api_extra/_agnostic/_manipulation.py b/src/array_api_extra/_agnostic/_manipulation.py new file mode 100644 index 00000000..50f6fcf9 --- /dev/null +++ b/src/array_api_extra/_agnostic/_manipulation.py @@ -0,0 +1,100 @@ +"""Array-agnostic implementations for manipulation functions.""" + +import math +from collections.abc import Sequence +from typing import cast + +from .._at import at +from .._lib import _compat +from .._lib._helpers import eager_shape, normalize_pad_width +from .._lib._typing import Array, ArrayNamespace + +__all__ = ["atleast_nd", "broadcast_shapes", "expand_dims", "pad"] + + +def atleast_nd(x: Array, /, *, ndim: int, xp: ArrayNamespace) -> Array: + # numpydoc ignore=PR01,RT01 + """See docstring in array_api_extra._delegation.""" + + if x.ndim < ndim: + x = xp.expand_dims(x, axis=0) + x = atleast_nd(x, ndim=ndim, xp=xp) + return x + + +# `float` in signature to accept `math.nan` for Dask. +# `int`s are still accepted as `float` is a superclass of `int` in typing +def broadcast_shapes( # numpydoc ignore=PR01,RT01 + *shapes: tuple[float | None, ...], +) -> tuple[int | None, ...]: + """See docstring in array_api_extra._delegation.""" + if not shapes: + return () # Match NumPy output + + ndim = max(len(shape) for shape in shapes) + out: list[int | None] = [] + for axis in range(-ndim, 0): + sizes = {shape[axis] for shape in shapes if axis >= -len(shape)} + # Dask uses NaN for unknown shape, which predates the Array API spec for None + none_size = None in sizes or math.nan in sizes # noqa: PLW0177 + sizes -= {1, None, math.nan} + if len(sizes) > 1: + msg = ( + "shape mismatch: objects cannot be broadcast to a single shape: " + f"{shapes}." + ) + raise ValueError(msg) + out.append(None if none_size else cast(int, sizes.pop()) if sizes else 1) + + return tuple(out) + + +def expand_dims( + a: Array, /, *, axis: tuple[int, ...] = (0,), xp: ArrayNamespace +) -> Array: + # numpydoc ignore=PR01,RT01 + """See docstring in array_api_extra._delegation.""" + for i in sorted(axis): + a = xp.expand_dims(a, axis=i) + return a + + +def pad( + x: Array, + pad_width: int | tuple[int, int] | Sequence[tuple[int, int]], + *, + constant_values: complex = 0, + xp: ArrayNamespace, +) -> Array: # numpydoc ignore=PR01,RT01 + """See docstring in `array_api_extra._delegation.py`.""" + pad_width_seq = normalize_pad_width(pad_width, x.ndim) + + slices: list[slice] = [] + newshape: list[int] = [] + for ax, w_tpl in enumerate(pad_width_seq): + if len(w_tpl) != 2: + msg = f"expect a 2-tuple (before, after), got {w_tpl}." + raise ValueError(msg) + + sh = eager_shape(x)[ax] + + if w_tpl[0] == 0 and w_tpl[1] == 0: + sl = slice(None, None, None) + else: + stop: int | None + start, stop = w_tpl + stop = None if stop == 0 else -stop + + sl = slice(start, stop, None) + sh += w_tpl[0] + w_tpl[1] + + newshape.append(sh) + slices.append(sl) + + padded = xp.full( + tuple(newshape), + fill_value=constant_values, + dtype=x.dtype, + device=_compat.device(x), + ) + return at(padded, tuple(slices)).set(x) diff --git a/src/array_api_extra/_agnostic/_searching.py b/src/array_api_extra/_agnostic/_searching.py new file mode 100644 index 00000000..b7b501ac --- /dev/null +++ b/src/array_api_extra/_agnostic/_searching.py @@ -0,0 +1,44 @@ +"""Array-agnostic implementations for searching functions.""" + +import math +from typing import Literal + +from .._lib import _compat +from .._lib._typing import Array, ArrayNamespace +from ._inspection import default_dtype + +__all__ = ["searchsorted"] + + +def searchsorted( + x1: Array, + x2: Array, + /, + *, + side: Literal["left", "right"] = "left", + xp: ArrayNamespace, +) -> Array: + # numpydoc ignore=PR01,RT01 + """See docstring in `array_api_extra._delegation.py`.""" + a = xp.full(x2.shape, 0, device=_compat.device(x1)) + + if x1.shape[-1] == 0: + return a + + n = xp.count_nonzero(~xp.isnan(x1), axis=-1, keepdims=True) + b = xp.broadcast_to(n, x2.shape) + + compare = xp.less_equal if side == "left" else xp.less + + # while xp.any(b - a > 1): + # refactored to for loop with ~log2(n) iterations for JAX JIT + for _ in range(int(math.log2(x1.shape[-1])) + 1): # type: ignore[arg-type] # pyright: ignore[reportArgumentType] + c = (a + b) // 2 + x0 = xp.take_along_axis(x1, c, axis=-1) + j = compare(x2, x0) + b = xp.where(j, c, b) + a = xp.where(j, a, c) + + out = xp.where(compare(x2, xp.min(x1, axis=-1, keepdims=True)), 0, b) + out = xp.where(xp.isnan(x2), x1.shape[-1], out) if side == "right" else out + return xp.astype(out, default_dtype(xp, kind="integral"), copy=False) diff --git a/src/array_api_extra/_agnostic/_set.py b/src/array_api_extra/_agnostic/_set.py new file mode 100644 index 00000000..2c7d5fa2 --- /dev/null +++ b/src/array_api_extra/_agnostic/_set.py @@ -0,0 +1,88 @@ +"""Array-agnostic implementations for set functions.""" + +from .._lib import _compat, _helpers +from .._lib._helpers import asarrays, capabilities +from .._lib._typing import Array, ArrayNamespace +from ._inspection import default_dtype + +__all__ = ["isin", "nunique", "setdiff1d", "union1d"] + + +def isin( # numpydoc ignore=PR01,RT01 + a: Array, + b: Array, + /, + *, + assume_unique: bool = False, + invert: bool = False, + xp: ArrayNamespace, +) -> Array: + """See docstring in `array_api_extra._delegation.py`.""" + original_a_shape = a.shape + a = xp.reshape(a, (-1,)) + b = xp.reshape(b, (-1,)) + return xp.reshape( + _helpers.in1d(a, b, assume_unique=assume_unique, invert=invert, xp=xp), + original_a_shape, + ) + + +def nunique(x: Array, /, *, xp: ArrayNamespace) -> Array: # numpydoc ignore=PR01,RT01 + """See docstring in `array_api_extra._delegation.py`.""" + # There are 3 general use cases: + # 1. backend has unique_counts and it returns an array with known shape + # 2. backend has unique_counts and it returns a None-sized array; + # e.g. Dask, ndonnx + # 3. backend does not have unique_counts; e.g. wrapped JAX + if capabilities(xp, device=_compat.device(x))["data-dependent shapes"]: + # xp has unique_counts; O(n) complexity + _, counts = xp.unique_counts(x) + n = _compat.size(counts) + if n is None: + return xp.sum(xp.ones_like(counts)) + return xp.asarray(n, device=_compat.device(x)) + + # xp does not have unique_counts; O(n*logn) complexity + x = xp.reshape(x, (-1,)) + x = xp.sort(x, stable=False) + mask = x != xp.roll(x, -1) + default_int = default_dtype(xp, "integral", device=_compat.device(x)) + return xp.maximum( + # Special cases: + # - array is size 0 + # - array has all elements equal to each other + xp.astype(xp.any(~mask), default_int), + xp.sum(xp.astype(mask, default_int)), + ) + + +def setdiff1d( + x1: Array | complex, + x2: Array | complex, + /, + *, + assume_unique: bool = False, + xp: ArrayNamespace, +) -> Array: # numpydoc ignore=PR01,RT01 + """See docstring in `array_api_extra._delegation.py`.""" + + # https://github.com/microsoft/pyright/issues/10103 + x1_, x2_ = asarrays(x1, x2, xp=xp) + + if assume_unique: + x1_ = xp.reshape(x1_, (-1,)) + x2_ = xp.reshape(x2_, (-1,)) + else: + x1_ = xp.unique_values(x1_) + x2_ = xp.unique_values(x2_) + + return x1_[_helpers.in1d(x1_, x2_, assume_unique=True, invert=True, xp=xp)] + + +def union1d(a: Array, b: Array, /, *, xp: ArrayNamespace) -> Array: + # numpydoc ignore=PR01,RT01 + """See docstring in `array_api_extra._delegation.py`.""" + a = xp.reshape(a, (-1,)) + b = xp.reshape(b, (-1,)) + # XXX: `sparse` returns NumPy arrays from `unique_values` + return xp.asarray(xp.unique_values(xp.concat([a, b]))) diff --git a/src/array_api_extra/_agnostic/_sorting.py b/src/array_api_extra/_agnostic/_sorting.py new file mode 100644 index 00000000..ecc1ca9b --- /dev/null +++ b/src/array_api_extra/_agnostic/_sorting.py @@ -0,0 +1,29 @@ +"""Array-agnostic implementations for sorting functions.""" + +from .._lib._typing import Array, ArrayNamespace + +__all__ = ["argpartition", "partition"] + + +def partition( # numpydoc ignore=PR01,RT01 + x: Array, + kth: int, # noqa: ARG001 + /, + axis: int = -1, + *, + xp: ArrayNamespace, +) -> Array: + """See docstring in `array_api_extra._delegation.py`.""" + return xp.sort(x, axis=axis, stable=False) + + +def argpartition( # numpydoc ignore=PR01,RT01 + x: Array, + kth: int, # noqa: ARG001 + /, + axis: int = -1, + *, + xp: ArrayNamespace, +) -> Array: + """See docstring in `array_api_extra._delegation.py`.""" + return xp.argsort(x, axis=axis, stable=False) diff --git a/src/array_api_extra/_agnostic/_statistical.py b/src/array_api_extra/_agnostic/_statistical.py new file mode 100644 index 00000000..02b39794 --- /dev/null +++ b/src/array_api_extra/_agnostic/_statistical.py @@ -0,0 +1,167 @@ +"""Array-agnostic implementations for statistical functions.""" + +import math +import warnings +from typing import cast + +from .._lib import _compat +from .._lib._helpers import eager_shape +from .._lib._typing import Array, ArrayNamespace +from ._manipulation import atleast_nd + +__all__ = ["cov", "nanmax", "nanmin", "nansum"] + + +def cov( + m: Array, + /, + *, + correction: float = 1, + fweights: Array | None = None, + aweights: Array | None = None, + xp: ArrayNamespace, +) -> Array: # numpydoc ignore=PR01,RT01 + """See docstring in array_api_extra._delegation.""" + # NB: no `xp.asarray(m)` here. The delegation layer already guarantees `m` + # is an array (it calls `array_namespace(m)` and reads `m.ndim`), and on + # torch `xp.asarray` detaches gradients and mutates the caller's tensor. + dtype = ( + xp.float64 if xp.isdtype(m.dtype, "integral") else xp.result_type(m, xp.float64) + ) + + m = atleast_nd(m, ndim=2, xp=xp) + # Preserve the historical no-alias guarantee even when the dtype already matches. + m = xp.astype(m, dtype, copy=True) + + # Validate weight shapes (eager metadata, lazy-safe). + n_obs = m.shape[-1] + for name, w_in in (("fweights", fweights), ("aweights", aweights)): + if w_in is None: + continue + if w_in.ndim != 1: + msg = f"`{name}` must be 1-D, got ndim={w_in.ndim}" + raise ValueError(msg) + weight_length = w_in.shape[0] + # Unknown dims are `None` per the standard; Dask non-standardly + # reports them as NaN, hence the `isnan` checks below. + if ( + weight_length is not None + and n_obs is not None + and not math.isnan(weight_length) + and not math.isnan(n_obs) + and weight_length != n_obs + ): + msg = ( + f"`{name}` has length {weight_length} but `m` has {n_obs} observations" + ) + raise ValueError(msg) + + fw = None + if fweights is not None: + fw = xp.astype(xp.asarray(fweights), dtype) + aw = None + if aweights is not None: + aw = xp.astype(xp.asarray(aweights), dtype) + if fw is None and aw is None: + w = None + elif fw is None: + w = aw + elif aw is None: + w = fw + else: + w = fw * aw + + if w is None: + avg = xp.mean(m, axis=-1, keepdims=True) + fact = eager_shape(m, axis=-1)[0] - correction + else: + v1 = xp.sum(w, axis=-1) + avg = xp.sum(m * w, axis=-1, keepdims=True) / v1 + if aw is None: + fact = v1 - correction + else: + fact = v1 - correction * xp.sum(w * aw, axis=-1) / v1 + + if not _compat.is_lazy_array(fact): + # Weights are cast to `dtype`, so a complex input produces a complex + # normalizer with a zero imaginary part. Complex ordering is undefined; + # compare its real component instead. + if w is not None: + fact_array = cast(Array, fact) + fact_to_check = ( + xp.real(fact_array) + if xp.isdtype(fact_array.dtype, "complex floating") + else fact_array + ) + else: + fact_to_check = fact + if fact_to_check <= 0: + warnings.warn( + "Degrees of freedom <= 0 for slice", RuntimeWarning, stacklevel=2 + ) + fact = 0 + + m_c = m - avg + m_w = m_c if w is None else m_c * w + m_cT = xp.matrix_transpose(m_c) + if xp.isdtype(m_cT.dtype, "complex floating"): + m_cT = xp.conj(m_cT) + c = m_w @ m_cT / fact + axes = tuple(axis for axis, length in enumerate(c.shape) if length == 1) + return xp.squeeze(c, axis=axes) + + +def nanmin( # numpydoc ignore=PR01,RT01 + a: Array, + /, + *, + axis: int | tuple[int, ...] | None, + xp: ArrayNamespace, +) -> Array: + """See docstring in `array_api_extra._delegation.py`.""" + mask = xp.isnan(a) + device_a = _compat.device(a) + x = xp.min( + xp.where(mask, xp.asarray(+xp.inf, dtype=a.dtype, device=device_a), a), + axis=axis, + ) + # Replace Infs from all NaN slices with NaN again + mask = xp.all(mask, axis=axis) + if xp.any(mask): + x = xp.where(mask, xp.asarray(xp.nan, dtype=x.dtype, device=device_a), x) + return x + + +def nanmax( # numpydoc ignore=PR01,RT01 + a: Array, + /, + *, + axis: int | tuple[int, ...] | None, + xp: ArrayNamespace, +) -> Array: + """See docstring in `array_api_extra._delegation.py`.""" + mask = xp.isnan(a) + device_a = _compat.device(a) + x = xp.max( + xp.where(mask, xp.asarray(-xp.inf, dtype=a.dtype, device=device_a), a), + axis=axis, + ) + # Replace Infs from all NaN slices with NaN again + mask = xp.all(mask, axis=axis) + if xp.any(mask): + x = xp.where(mask, xp.asarray(xp.nan, dtype=x.dtype, device=device_a), x) + return x + + +def nansum( # numpydoc ignore=PR01,RT01 + a: Array, + /, + *, + axis: int | tuple[int, ...] | None, + xp: ArrayNamespace, +) -> Array: + """See docstring in `array_api_extra._delegation.py`.""" + mask = xp.isnan(a) + device_a = _compat.device(a) + zero = xp.asarray(0, dtype=a.dtype, device=device_a) + return xp.sum(xp.where(mask, zero, a), axis=axis) diff --git a/src/array_api_extra/_lib/_at.py b/src/array_api_extra/_at.py similarity index 97% rename from src/array_api_extra/_lib/_at.py rename to src/array_api_extra/_at.py index 018ea2a8..9ef605d5 100644 --- a/src/array_api_extra/_lib/_at.py +++ b/src/array_api_extra/_at.py @@ -7,21 +7,23 @@ from enum import Enum from typing import TYPE_CHECKING, ClassVar, cast -from ._utils import _compat -from ._utils._compat import ( +from ._lib import _compat +from ._lib._compat import ( array_namespace, is_dask_array, is_jax_array, is_torch_array, is_writeable_array, ) -from ._utils._helpers import meta_namespace -from ._utils._typing import Array, ArrayNamespace, SetIndex +from ._lib._helpers import meta_namespace +from ._lib._typing import Array, ArrayNamespace, SetIndex if TYPE_CHECKING: # pragma: no cover # TODO import from typing (requires Python >=3.11) from typing import Self +__all__ = ["at"] + class _AtOp(Enum): """Operations for use in `xpx.at`.""" @@ -48,13 +50,13 @@ def __str__(self) -> str: # pyright: ignore[reportImplicitOverride] return self.value -class Undef(Enum): +class _Undef(Enum): """Sentinel for undefined values.""" UNDEF = 0 -_undef = Undef.UNDEF +_undef = _Undef.UNDEF class at: # pylint: disable=invalid-name # numpydoc ignore=PR02 @@ -199,11 +201,11 @@ class at: # pylint: disable=invalid-name # numpydoc ignore=PR02 """ _x: Array - _idx: SetIndex | Undef + _idx: SetIndex | _Undef __slots__: ClassVar[tuple[str, ...]] = ("_idx", "_x") def __init__( - self, x: Array, idx: SetIndex | Undef = _undef, / + self, x: Array, idx: SetIndex | _Undef = _undef, / ) -> None: # numpydoc ignore=GL08 self._x = x self._idx = idx @@ -268,12 +270,12 @@ def _op( Array Updated `x`. """ - from ._funcs import apply_where # pylint: disable=cyclic-import + from ._agnostic._elementwise import apply_where # pylint: disable=cyclic-import x, idx = self._x, self._idx xp = array_namespace(x, y) if xp is None else xp - if isinstance(idx, Undef): + if isinstance(idx, _Undef): msg = ( "Index has not been set.\n" "Usage: either\n" diff --git a/src/array_api_extra/_creation.py b/src/array_api_extra/_creation.py new file mode 100644 index 00000000..eed8e016 --- /dev/null +++ b/src/array_api_extra/_creation.py @@ -0,0 +1,154 @@ +"""Delegation layer for creation functions.""" + +from ._agnostic import _creation, _inspection +from ._lib._compat import ( + array_namespace, + is_cupy_namespace, + is_dask_namespace, + is_jax_namespace, + is_numpy_namespace, + is_torch_namespace, +) +from ._lib._compat import device as get_device +from ._lib._typing import Array, ArrayNamespace, DType + +__all__ = ["create_diagonal", "one_hot"] + + +def create_diagonal( + x: Array, /, *, offset: int = 0, xp: ArrayNamespace | None = None +) -> Array: + """ + Construct a diagonal array. + + Parameters + ---------- + x : array + An array having shape ``(*batch_dims, k)``. + offset : int, optional + Offset from the leading diagonal (default is ``0``). + Use positive ints for diagonals above the leading diagonal, + and negative ints for diagonals below the leading diagonal. + xp : array_namespace, optional + The standard-compatible namespace for `x`. Default: infer. + + Returns + ------- + array + An array having shape ``(*batch_dims, k+abs(offset), k+abs(offset))`` with `x` + on the diagonal (offset by `offset`). + + Examples + -------- + >>> import array_api_strict as xp + >>> import array_api_extra as xpx + >>> x = xp.asarray([2, 4, 8]) + + >>> xpx.create_diagonal(x, xp=xp) + Array([[2, 0, 0], + [0, 4, 0], + [0, 0, 8]], dtype=array_api_strict.int64) + + >>> xpx.create_diagonal(x, offset=-2, xp=xp) + Array([[0, 0, 0, 0, 0], + [0, 0, 0, 0, 0], + [2, 0, 0, 0, 0], + [0, 4, 0, 0, 0], + [0, 0, 8, 0, 0]], dtype=array_api_strict.int64) + """ + if xp is None: + xp = array_namespace(x) + + if x.ndim == 0: + err_msg = "`x` must be at least 1-dimensional." + raise ValueError(err_msg) + + if is_torch_namespace(xp): + return xp.diag_embed(x, offset=offset, dim1=-2, dim2=-1) + + if ( + is_dask_namespace(xp) + or is_cupy_namespace(xp) + or is_numpy_namespace(xp) + or is_jax_namespace(xp) + ) and (x.ndim < 2): + return xp.diag(x, k=offset) + + return _creation.create_diagonal(x, offset=offset, xp=xp) + + +def one_hot( + x: Array, + /, + num_classes: int, + *, + dtype: DType | None = None, + axis: int = -1, + xp: ArrayNamespace | None = None, +) -> Array: + """ + One-hot encode the given indices. + + Each index in the input `x` is encoded as a vector of zeros of length `num_classes` + with the element at the given index set to one. + + Parameters + ---------- + x : array + An array with integral dtype whose values are between `0` and `num_classes - 1`. + num_classes : int + Number of classes in the one-hot dimension. + dtype : DType, optional + The dtype of the return value. Defaults to the default float dtype (usually + float64). + axis : int, optional + Position in the expanded axes where the new axis is placed. Default: -1. + xp : array_namespace, optional + The standard-compatible namespace for `x`. Default: infer. + + Returns + ------- + array + An array having the same shape as `x` except for a new axis at the position + given by `axis` having size `num_classes`. If `axis` is unspecified, it + defaults to -1, which appends a new axis. + + If ``x < 0`` or ``x >= num_classes``, then the result is undefined, may raise + an exception, or may even cause a bad state. `x` is not checked. + + Examples + -------- + >>> import array_api_extra as xpx + >>> import array_api_strict as xp + >>> xpx.one_hot(xp.asarray([1, 2, 0]), 3) + Array([[0., 1., 0.], + [0., 0., 1.], + [1., 0., 0.]], dtype=array_api_strict.float64) + """ + # Validate inputs. + if xp is None: + xp = array_namespace(x) + if not xp.isdtype(x.dtype, "integral"): + msg = "x must have an integral dtype." + raise TypeError(msg) + if dtype is None: + dtype = _inspection.default_dtype(xp, device=get_device(x)) + # Delegate where possible. + if is_jax_namespace(xp): + from jax.nn import one_hot as jax_one_hot + + return jax_one_hot(x, num_classes, dtype=dtype, axis=axis) + if is_torch_namespace(xp): + from torch.nn.functional import one_hot as torch_one_hot + + x = xp.astype(x, xp.int64) # PyTorch only supports int64 here. + try: + out = torch_one_hot(x, num_classes) + except RuntimeError as e: + raise IndexError from e + else: + out = _creation.one_hot(x, num_classes, xp=xp) + out = xp.astype(out, dtype, copy=False) + if axis != -1: + out = xp.moveaxis(out, -1, axis) + return out diff --git a/src/array_api_extra/_delegation.py b/src/array_api_extra/_delegation.py deleted file mode 100644 index c91517ca..00000000 --- a/src/array_api_extra/_delegation.py +++ /dev/null @@ -1,1998 +0,0 @@ -"""Delegation to existing implementations for Public API Functions.""" - -from collections.abc import Sequence -from typing import Literal - -from ._lib import _funcs -from ._lib._utils._compat import ( - array_namespace, - is_cupy_namespace, - is_dask_namespace, - is_jax_array, - is_jax_namespace, - is_numpy_namespace, - is_pydata_sparse_namespace, - is_torch_namespace, - size, -) -from ._lib._utils._compat import device as get_device -from ._lib._utils._helpers import ( - asarrays, - capabilities, - deprecated, - eager_shape, - normalize_pad_width, -) -from ._lib._utils._typing import Array, ArrayNamespace, Device, DType - -__all__ = [ - "argpartition", - "atleast_nd", - "broadcast_shapes", - "cov", - "create_diagonal", - "deg2rad", - "diag_indices", - "expand_dims", - "isclose", - "isin", - "kron", - "nan_to_num", - "nanmax", - "nanmin", - "nansum", - "nunique", - "one_hot", - "pad", - "partition", - "rad2deg", - "searchsorted", - "setdiff1d", - "sinc", - "tril_indices", - "triu_indices", - "union1d", - "unravel_index", -] - - -def atleast_nd(x: Array, /, *, ndim: int, xp: ArrayNamespace | None = None) -> Array: - """ - Recursively expand the dimension of an array to at least `ndim`. - - Parameters - ---------- - x : array - Input array. - ndim : int - The minimum number of dimensions for the result. - xp : array_namespace, optional - The standard-compatible namespace for `x`. Default: infer. - - Returns - ------- - array - An array with ``res.ndim`` >= `ndim`. - If ``x.ndim`` >= `ndim`, `x` is returned. - If ``x.ndim`` < `ndim`, `x` is expanded by prepending new axes - until ``res.ndim`` equals `ndim`. - - Examples - -------- - >>> import array_api_strict as xp - >>> import array_api_extra as xpx - >>> x = xp.asarray([1]) - >>> xpx.atleast_nd(x, ndim=3, xp=xp) - Array([[[1]]], dtype=array_api_strict.int64) - - >>> x = xp.asarray([[[1, 2], - ... [3, 4]]]) - >>> xpx.atleast_nd(x, ndim=1, xp=xp) is x - True - """ - if xp is None: - xp = array_namespace(x) - - if 1 <= ndim <= 2 and ( - is_numpy_namespace(xp) - or is_jax_namespace(xp) - or is_dask_namespace(xp) - or is_cupy_namespace(xp) - or is_torch_namespace(xp) - ): - return getattr(xp, f"atleast_{ndim}d")(x) - - return _funcs.atleast_nd(x, ndim=ndim, xp=xp) - - -@deprecated( - "`xpx.broadcast_shapes` is deprecated and will be removed in v1.0.0. " - "`xp.broadcast_shapes` exists in the standard as of v2025.12." -) -def broadcast_shapes( - *shapes: tuple[float | None, ...], xp: ArrayNamespace | None = None -) -> tuple[int | None, ...]: - """ - Compute the shape of the broadcasted arrays. - - .. deprecated:: 0.11.0 - :func:`broadcast_shapes` is deprecated and will be removed in v1.0.0. - :func:`array_api.broadcast_shapes` exists in the standard as of v2025.12. - - Duplicates :func:`numpy.broadcast_shapes`, with additional support for - None and NaN sizes. - - Parameters - ---------- - *shapes : tuple[int | None, ...] - Shapes of the arrays to broadcast. - xp : array_namespace, optional - The standard-compatible namespace to use for native delegation. - Default: use the array-agnostic implementation. - - Returns - ------- - tuple[int | None, ...] - The shape of the broadcasted arrays. - - See Also - -------- - numpy.broadcast_shapes : Equivalent NumPy function. - array_api.broadcast_arrays : Function to broadcast actual arrays. - - Notes - ----- - This function accepts the Array API's ``None`` for unknown sizes, - as well as Dask's non-standard ``math.nan``. - Regardless of input, the output always contains ``None`` for unknown sizes. - - Examples - -------- - >>> import array_api_extra as xpx - >>> xpx.broadcast_shapes((2, 3), (2, 1)) - (2, 3) - >>> xpx.broadcast_shapes((4, 2, 3), (2, 1), (1, 3)) - (4, 2, 3) - """ - if ( - xp is not None - and all(isinstance(size, int) for shape in shapes for size in shape) - and ( - is_numpy_namespace(xp) - or is_cupy_namespace(xp) - or is_jax_namespace(xp) - or is_torch_namespace(xp) - ) - ): - return xp.broadcast_shapes(*shapes) - - return _funcs.broadcast_shapes(*shapes) - - -def cov( - m: Array, - /, - *, - axis: int = -1, - correction: float = 1, - fweights: Array | None = None, - aweights: Array | None = None, - xp: ArrayNamespace | None = None, -) -> Array: - """ - Estimate a covariance matrix (or a stack of covariance matrices). - - Covariance indicates the level to which two variables vary together. - If we examine *N*-dimensional samples, :math:`X = [x_1, x_2, ... x_N]^T`, - each with *M* observations, then element :math:`C_{ij}` of the - :math:`N \\times N` covariance matrix is the covariance of - :math:`x_i` and :math:`x_j`. The element :math:`C_{ii}` is the variance - of :math:`x_i`. - - Extends :func:`numpy.cov` with support for batch input. - Naming follows the array API conventions used elsewhere in - this library (``axis``, ``correction``) rather than the NumPy spellings - (``rowvar``, ``bias``, ``ddof``); see Notes for the mapping. - - Parameters - ---------- - m : array - An array of shape ``(..., N, M)`` whose innermost two dimensions - contain *M* observations of *N* variables by default. The axis of - observations is controlled by `axis`. - axis : int, optional - Axis of `m` containing the observations. Default: ``-1`` (the last - axis), matching the array API convention. Use ``axis=-2`` (or ``0`` - for 2-D input) to treat each column as a variable, which - corresponds to ``rowvar=False`` in :func:`numpy.cov`. - correction : int or float, optional - Degrees of freedom correction: normalization divides by - ``N - correction`` (for unweighted input). Default: ``1``, which - gives the unbiased estimate (matches :func:`numpy.cov` default of - ``bias=False``). Set to ``0`` for the biased estimate (``N`` - normalization). Corresponds to ``ddof`` in :func:`numpy.cov` and to - ``correction`` in :func:`numpy.var`/:func:`numpy.std` and - :func:`torch.cov`. - Non-integer values are allowed for advanced use cases: the - unbiased correction for weighted observations depends on the - sum and dispersion of the weights and is generally not an - integer, and autocorrelated data may also require a fractional - correction. Non-integer ``correction`` routes through the - generic implementation because :func:`numpy.cov`'s ``ddof`` and - :func:`torch.cov`'s ``correction`` both require integers. - fweights : array, optional - 1-D array of integer frequency weights: the number of times each - observation is repeated. Same as ``fweights`` in - :func:`numpy.cov`/:func:`torch.cov`. - aweights : array, optional - 1-D array of observation-vector weights (analytic weights). Larger - values mark more important observations. Same as ``aweights`` in - :func:`numpy.cov`/:func:`torch.cov`. - xp : array_namespace, optional - The standard-compatible namespace for `m`. Default: infer. - - Returns - ------- - array - An array having shape ``(..., N, N)`` whose innermost two dimensions represent - the covariance matrix of the variables. - - Notes - ----- - Mapping from :func:`numpy.cov` to this function:: - - numpy.cov(m, rowvar=True) -> cov(m, axis=-1) # default - numpy.cov(m, rowvar=False) -> cov(m, axis=-2) - numpy.cov(m, bias=True) -> cov(m, correction=0) - numpy.cov(m, ddof=k) -> cov(m, correction=k) - numpy.cov(m, fweights=f) -> cov(m, fweights=f) - numpy.cov(m, aweights=a) -> cov(m, aweights=a) - - A ``RuntimeWarning`` is emitted for non-positive effective degrees of - freedom when the effective normalizer can be checked without materializing - a lazy array. When the normalizer itself is lazy (e.g. for weighted Dask - inputs), this check is skipped; choose ``correction`` and weights such that - it is positive. - - Examples - -------- - >>> import array_api_strict as xp - >>> import array_api_extra as xpx - - Consider two variables, :math:`x_0` and :math:`x_1`, which - correlate perfectly, but in opposite directions: - - >>> x = xp.asarray([[0, 2], [1, 1], [2, 0]]).T - >>> x - Array([[0, 1, 2], - [2, 1, 0]], dtype=array_api_strict.int64) - - Note how :math:`x_0` increases while :math:`x_1` decreases. The covariance - matrix shows this clearly: - - >>> xpx.cov(x, xp=xp) - Array([[ 1., -1.], - [-1., 1.]], dtype=array_api_strict.float64) - - Note that element :math:`C_{0,1}`, which shows the correlation between - :math:`x_0` and :math:`x_1`, is negative. - - Further, note how `x` and `y` are combined: - - >>> x = xp.asarray([-2.1, -1, 4.3]) - >>> y = xp.asarray([3, 1.1, 0.12]) - >>> X = xp.stack((x, y), axis=0) - >>> xpx.cov(X, xp=xp) - Array([[11.71 , -4.286 ], - [-4.286 , 2.14413333]], dtype=array_api_strict.float64) - - >>> xpx.cov(x, xp=xp) - Array(11.71, dtype=array_api_strict.float64) - - >>> xpx.cov(y, xp=xp) - Array(2.14413333, dtype=array_api_strict.float64) - - Input with more than two dimensions is treated as a stack of - two-dimensional input. - - >>> stack = xp.stack((X, 2*X)) - >>> xpx.cov(stack) - Array([[[ 11.71 , -4.286 ], - [ -4.286 , 2.14413333]], - [[ 46.84 , -17.144 ], - [-17.144 , 8.57653333]]], dtype=array_api_strict.float64) - - The normalization can be adjusted with `correction`, and observations - can be weighted with integer frequencies `fweights` or importance - weights `aweights`: - - >>> x = xp.asarray([0., 1., 2., 3., 4.]) - >>> xpx.cov(x, xp=xp) # unbiased variance: divide by N - 1 - Array(2.5, dtype=array_api_strict.float64) - >>> xpx.cov(x, correction=0, xp=xp) # biased variance: divide by N - Array(2., dtype=array_api_strict.float64) - - Giving the two extreme observations frequency 2 via `fweights` is - equivalent to repeating them in `x`: - - >>> xpx.cov(x, fweights=xp.asarray([2, 1, 1, 1, 2]), xp=xp) - Array(3., dtype=array_api_strict.float64) - >>> xpx.cov(xp.asarray([0., 0., 1., 2., 3., 4., 4.]), xp=xp) - Array(3., dtype=array_api_strict.float64) - - `aweights` instead adjusts the relative importance of observations, - here down-weighting the two extremes: - - >>> xpx.cov(x, aweights=xp.asarray([0.5, 1., 1., 1., 0.5]), xp=xp) - Array(1.92, dtype=array_api_strict.float64) - """ - - if xp is None: - xp = array_namespace(m, fweights, aweights) - - # Validate axis against m.ndim. - ndim = max(m.ndim, 1) - if not -ndim <= axis < ndim: - msg = f"axis {axis} is out of bounds for array of dimension {m.ndim}" - raise IndexError(msg) - - # Normalize: observations on the last axis. After this, every backend - # sees the same convention and we never need to deal with `rowvar`. - if m.ndim >= 2 and axis not in (-1, m.ndim - 1): - m = xp.moveaxis(m, axis, -1) - - # `numpy.cov` (and cupy/dask/jax) require integer `ddof`; `torch.cov` - # requires integer `correction`. For non-integer-valued `correction`, - # fall through to the generic implementation. - integer_correction = float(correction).is_integer() - has_weights = fweights is not None or aweights is not None - - if m.ndim <= 2 and integer_correction: - # Not just for static typing: `correction` may be an integer-valued - # float such as 1.0, which `torch.cov` rejects at runtime. - int_correction = int(correction) - if is_torch_namespace(xp): - fw = None if fweights is None else xp.asarray(fweights) - aw = None if aweights is None else xp.asarray(aweights) - return xp.cov(m, correction=int_correction, fweights=fw, aweights=aw) - # `dask.array.cov` forces `.compute()` whenever weights are given: - # its internal `if fact <= 0` check on a lazy 0-D scalar triggers - # materialization. Route to the generic impl, which is fully lazy - # because it only does sum/matmul and skips that scalar check. - if ( - is_numpy_namespace(xp) - or is_cupy_namespace(xp) - or is_jax_namespace(xp) - or (is_dask_namespace(xp) and not has_weights) - ): - return xp.cov( - m, - ddof=int_correction, - fweights=fweights, - aweights=aweights, - ) - - return _funcs.cov( - m, - correction=correction, - fweights=fweights, - aweights=aweights, - xp=xp, - ) - - -def create_diagonal( - x: Array, /, *, offset: int = 0, xp: ArrayNamespace | None = None -) -> Array: - """ - Construct a diagonal array. - - Parameters - ---------- - x : array - An array having shape ``(*batch_dims, k)``. - offset : int, optional - Offset from the leading diagonal (default is ``0``). - Use positive ints for diagonals above the leading diagonal, - and negative ints for diagonals below the leading diagonal. - xp : array_namespace, optional - The standard-compatible namespace for `x`. Default: infer. - - Returns - ------- - array - An array having shape ``(*batch_dims, k+abs(offset), k+abs(offset))`` with `x` - on the diagonal (offset by `offset`). - - Examples - -------- - >>> import array_api_strict as xp - >>> import array_api_extra as xpx - >>> x = xp.asarray([2, 4, 8]) - - >>> xpx.create_diagonal(x, xp=xp) - Array([[2, 0, 0], - [0, 4, 0], - [0, 0, 8]], dtype=array_api_strict.int64) - - >>> xpx.create_diagonal(x, offset=-2, xp=xp) - Array([[0, 0, 0, 0, 0], - [0, 0, 0, 0, 0], - [2, 0, 0, 0, 0], - [0, 4, 0, 0, 0], - [0, 0, 8, 0, 0]], dtype=array_api_strict.int64) - """ - if xp is None: - xp = array_namespace(x) - - if x.ndim == 0: - err_msg = "`x` must be at least 1-dimensional." - raise ValueError(err_msg) - - if is_torch_namespace(xp): - return xp.diag_embed(x, offset=offset, dim1=-2, dim2=-1) - - if ( - is_dask_namespace(xp) - or is_cupy_namespace(xp) - or is_numpy_namespace(xp) - or is_jax_namespace(xp) - ) and (x.ndim < 2): - return xp.diag(x, k=offset) - - return _funcs.create_diagonal(x, offset=offset, xp=xp) - - -def deg2rad(x: Array, /, *, xp: ArrayNamespace | None = None) -> Array: - """ - Convert angles from degrees to radians. - - Parameters - ---------- - x : array - Input array in degrees. Must have an integral or floating-point dtype. - xp : array_namespace, optional - The standard-compatible namespace for `x`. Default: infer. - - Returns - ------- - array - The corresponding angles in radians. Integral inputs are converted to the - default floating-point dtype. - - Examples - -------- - >>> import array_api_strict as xp - >>> import array_api_extra as xpx - >>> xpx.deg2rad(xp.asarray([0, 90, 180]), xp=xp) - Array([0. , 1.57079633, 3.14159265], dtype=array_api_strict.float64) - """ - if xp is None: - xp = array_namespace(x) - if xp.isdtype(x.dtype, "integral"): - x = xp.astype(x, _funcs.default_dtype(xp, device=get_device(x))) - elif not xp.isdtype(x.dtype, ("real floating", "complex floating")): - msg = "`x` must have an integral, real floating, or complex floating dtype." - raise TypeError(msg) - - if is_jax_namespace(xp) or ( - not xp.isdtype(x.dtype, "complex floating") - and ( - is_numpy_namespace(xp) - or is_cupy_namespace(xp) - or is_torch_namespace(xp) - or is_dask_namespace(xp) - ) - ): - return xp.deg2rad(x) - - return _funcs.deg2rad(x, xp=xp) - - -def rad2deg(x: Array, /, *, xp: ArrayNamespace | None = None) -> Array: - """ - Convert angles from radians to degrees. - - Parameters - ---------- - x : array - Input array in radians. Must have an integral or floating-point dtype. - xp : array_namespace, optional - The standard-compatible namespace for `x`. Default: infer. - - Returns - ------- - array - The corresponding angles in degrees. Integral inputs are converted to the - default floating-point dtype. - - Examples - -------- - >>> import array_api_strict as xp - >>> import array_api_extra as xpx - >>> xpx.rad2deg(xp.asarray([0.0, xp.pi / 2, xp.pi]), xp=xp) - Array([ 0., 90., 180.], dtype=array_api_strict.float64) - """ - if xp is None: - xp = array_namespace(x) - if xp.isdtype(x.dtype, "integral"): - x = xp.astype(x, _funcs.default_dtype(xp, device=get_device(x))) - elif not xp.isdtype(x.dtype, ("real floating", "complex floating")): - msg = "`x` must have an integral, real floating, or complex floating dtype." - raise TypeError(msg) - - if is_jax_namespace(xp) or ( - not xp.isdtype(x.dtype, "complex floating") - and ( - is_numpy_namespace(xp) - or is_cupy_namespace(xp) - or is_torch_namespace(xp) - or is_dask_namespace(xp) - ) - ): - return xp.rad2deg(x) - - return _funcs.rad2deg(x, xp=xp) - - -def diag_indices( - n: int, /, *, ndim: int = 2, device: Device | None = None, xp: ArrayNamespace -) -> tuple[Array, ...]: - """ - Return the indices to access the main diagonal of an array. - - Equivalent to :func:`numpy.diag_indices`. - - Parameters - ---------- - n : int - The size of each dimension of the (hyper-)cube ``(n, n, ..., n)`` - that the returned indices index into. - ndim : int, optional - The number of dimensions. Default: ``2``. - device : Device, optional - The device on which to place the returned arrays. Default: current device. - xp : array_namespace - The standard-compatible namespace to create the indices in. - - Returns - ------- - tuple of array - 1-D integer arrays of length ``n`` that together index - the main diagonal of an array of shape ``(n,) * ndim``. - - Examples - -------- - >>> import array_api_strict as xp - >>> import array_api_extra as xpx - >>> rows, cols = xpx.diag_indices(3, xp=xp) - >>> rows - Array([0, 1, 2], dtype=array_api_strict.int64) - >>> cols - Array([0, 1, 2], dtype=array_api_strict.int64) - """ - if n < 0: - msg = f"`n` must be non-negative, got {n}" - raise ValueError(msg) - if ndim < 1: - msg = f"`ndim` must be >= 1, got {ndim}" - raise ValueError(msg) - if device is None and ( - is_numpy_namespace(xp) or is_cupy_namespace(xp) or is_jax_namespace(xp) - ): - return xp.diag_indices(n, ndim=ndim) - return _funcs.diag_indices(n, ndim=ndim, device=device, xp=xp) - - -@deprecated( - "`xpx.expand_dims` is deprecated and will be removed in v1.0.0. " - "`xp.expand_dims` with support for a tuple of ints in `axis` " - "exists in the standard as of v2025.12." -) -def expand_dims( - a: Array, /, *, axis: int | tuple[int, ...] = (0,), xp: ArrayNamespace | None = None -) -> Array: - """ - Expand the shape of an array. - - .. deprecated:: 0.11.0 - :func:`expand_dims` is deprecated and will be removed in v1.0.0. - :func:`array_api.expand_dims` with support for a tuple of ints in `axis` - exists in the standard as of v2025.12. - - Insert (a) new axis/axes that will appear at the position(s) specified by - `axis` in the expanded array shape. - - Parameters - ---------- - a : array - Array to have its shape expanded. - axis : int or tuple of ints, optional - Position(s) in the expanded axes where the new axis (or axes) is/are placed. - If multiple positions are provided, they should be unique (note that a position - given by a positive index could also be referred to by a negative index - - that will also result in an error). - Default: ``(0,)``. - xp : array_namespace, optional - The standard-compatible namespace for `a`. Default: infer. - - Returns - ------- - array - `a` with an expanded shape. - - Examples - -------- - >>> import array_api_strict as xp - >>> import array_api_extra as xpx - >>> x = xp.asarray([1, 2]) - >>> x.shape - (2,) - - The following is equivalent to ``x[xp.newaxis, :]`` or ``x[xp.newaxis]``: - - >>> y = xpx.expand_dims(x, axis=0, xp=xp) - >>> y - Array([[1, 2]], dtype=array_api_strict.int64) - >>> y.shape - (1, 2) - - The following is equivalent to ``x[:, xp.newaxis]``: - - >>> y = xpx.expand_dims(x, axis=1, xp=xp) - >>> y - Array([[1], - [2]], dtype=array_api_strict.int64) - >>> y.shape - (2, 1) - - ``axis`` may also be a tuple: - - >>> y = xpx.expand_dims(x, axis=(0, 1), xp=xp) - >>> y - Array([[[1, 2]]], dtype=array_api_strict.int64) - - >>> y = xpx.expand_dims(x, axis=(2, 0), xp=xp) - >>> y - Array([[[1], - [2]]], dtype=array_api_strict.int64) - """ - if xp is None: - xp = array_namespace(a) - - if not isinstance(axis, tuple): - axis = (axis,) - ndim = a.ndim + len(axis) - if axis != () and (min(axis) < -ndim or max(axis) >= ndim): - err_msg = ( - f"a provided axis position is out of bounds for array of dimension {a.ndim}" - ) - raise IndexError(err_msg) - axis = tuple(dim % ndim for dim in axis) - if len(set(axis)) != len(axis): - err_msg = "Duplicate dimensions specified in `axis`." - raise ValueError(err_msg) - - if is_numpy_namespace(xp) or is_dask_namespace(xp) or is_jax_namespace(xp): - return xp.expand_dims(a, axis=axis) - - return _funcs.expand_dims(a, axis=axis, xp=xp) - - -def isclose( - a: Array | complex, - b: Array | complex, - *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, - xp: ArrayNamespace | None = None, -) -> Array: - """ - Return a boolean array where two arrays are element-wise equal within a tolerance. - - The tolerance values are positive, typically very small numbers. The relative - difference ``(rtol * abs(b))`` and the absolute difference `atol` are added together - to compare against the absolute difference between `a` and `b`. - - NaNs are treated as equal if they are in the same place and if ``equal_nan=True``. - Infs are treated as equal if they are in the same place and of the same sign in both - arrays. - - Parameters - ---------- - a, b : Array | int | float | complex | bool - Input objects to compare. At least one must be an array. - rtol : array_like, optional - The relative tolerance parameter (see Notes). - atol : array_like, optional - The absolute tolerance parameter (see Notes). - equal_nan : bool, optional - Whether to compare NaN's as equal. If True, NaN's in `a` will be considered - equal to NaN's in `b` in the output array. - xp : array_namespace, optional - The standard-compatible namespace for `a` and `b`. Default: infer. - - Returns - ------- - Array - A boolean array of shape broadcasted from `a` and `b`, containing ``True`` where - `a` is close to `b`, and ``False`` otherwise. - - Warnings - -------- - The default `atol` is not appropriate for comparing numbers with magnitudes much - smaller than one (see notes). - - See Also - -------- - math.isclose : Similar function in stdlib for Python scalars. - - Notes - ----- - For finite values, `isclose` uses the following equation to test whether two - floating point values are equivalent:: - - absolute(a - b) <= (atol + rtol * absolute(b)) - - Unlike the built-in `math.isclose`, - the above equation is not symmetric in `a` and `b`, - so that ``isclose(a, b)`` might be different from ``isclose(b, a)`` in some rare - cases. - - The default value of `atol` is not appropriate when the reference value `b` has - magnitude smaller than one. For example, it is unlikely that ``a = 1e-9`` and - ``b = 2e-9`` should be considered "close", yet ``isclose(1e-9, 2e-9)`` is ``True`` - with default settings. Be sure to select `atol` for the use case at hand, especially - for defining the threshold below which a non-zero value in `a` will be considered - "close" to a very small or zero value in `b`. - - The comparison of `a` and `b` uses standard broadcasting, which means that `a` and - `b` need not have the same shape in order for ``isclose(a, b)`` to evaluate to - ``True``. - - `isclose` is not defined for non-numeric data types. - ``bool`` is considered a numeric data-type for this purpose. - """ - xp = array_namespace(a, b) if xp is None else xp - - if ( - is_numpy_namespace(xp) - or is_cupy_namespace(xp) - or is_dask_namespace(xp) - or is_jax_namespace(xp) - ): - return xp.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) - - if is_torch_namespace(xp): - a, b = asarrays(a, b, xp=xp) # Array API 2024.12 support - return xp.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) - - return _funcs.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan, xp=xp) - - -def kron( - a: Array | complex, - b: Array | complex, - /, - *, - xp: ArrayNamespace | None = None, -) -> Array: - """ - Kronecker product of two arrays. - - Computes the Kronecker product, a composite array made of blocks of the - second array scaled by the first. - - Equivalent to ``numpy.kron`` for NumPy arrays. - - Parameters - ---------- - a, b : Array | int | float | complex - Input arrays or scalars. At least one must be an array. - xp : array_namespace, optional - The standard-compatible namespace for `a` and `b`. Default: infer. - - Returns - ------- - array - The Kronecker product of `a` and `b`. - - Notes - ----- - The function assumes that the number of dimensions of `a` and `b` - are the same, if necessary prepending the smallest with ones. - If ``a.shape = (r0,r1,..,rN)`` and ``b.shape = (s0,s1,...,sN)``, - the Kronecker product has shape ``(r0*s0, r1*s1, ..., rN*SN)``. - The elements are products of elements from `a` and `b`, organized - explicitly by:: - - kron(a,b)[k0,k1,...,kN] = a[i0,i1,...,iN] * b[j0,j1,...,jN] - - where:: - - kt = it * st + jt, t = 0,...,N - - In the common 2-D case (N=1), the block structure can be visualized:: - - [[ a[0,0]*b, a[0,1]*b, ... , a[0,-1]*b ], - [ ... ... ], - [ a[-1,0]*b, a[-1,1]*b, ... , a[-1,-1]*b ]] - - Examples - -------- - >>> import array_api_strict as xp - >>> import array_api_extra as xpx - >>> xpx.kron(xp.asarray([1, 10, 100]), xp.asarray([5, 6, 7]), xp=xp) - Array([ 5, 6, 7, 50, 60, 70, 500, - 600, 700], dtype=array_api_strict.int64) - - >>> xpx.kron(xp.asarray([5, 6, 7]), xp.asarray([1, 10, 100]), xp=xp) - Array([ 5, 50, 500, 6, 60, 600, 7, - 70, 700], dtype=array_api_strict.int64) - - >>> xpx.kron(xp.eye(2), xp.ones((2, 2)), xp=xp) - Array([[1., 1., 0., 0.], - [1., 1., 0., 0.], - [0., 0., 1., 1.], - [0., 0., 1., 1.]], dtype=array_api_strict.float64) - - >>> a = xp.reshape(xp.arange(100), (2, 5, 2, 5)) - >>> b = xp.reshape(xp.arange(24), (2, 3, 4)) - >>> c = xpx.kron(a, b, xp=xp) - >>> c.shape - (2, 10, 6, 20) - >>> I = (1, 3, 0, 2) - >>> J = (0, 2, 1) - >>> J1 = (0,) + J # extend to ndim=4 - >>> S1 = (1,) + b.shape - >>> K = tuple(xp.asarray(I) * xp.asarray(S1) + xp.asarray(J1)) - >>> c[K] == a[I]*b[J] - Array(True, dtype=array_api_strict.bool) - """ - if xp is None: - xp = array_namespace(a, b) - - a, b = asarrays(a, b, xp=xp) - - if ( - is_cupy_namespace(xp) - or is_jax_namespace(xp) - or is_numpy_namespace(xp) - or is_torch_namespace(xp) - ): - return xp.kron(a, b) - - return _funcs.kron(a, b, xp=xp) - - -def nan_to_num( - x: Array | complex, - /, - *, - fill_value: float = 0.0, - xp: ArrayNamespace | None = None, -) -> Array: - """ - Replace NaN with zero and infinity with large finite numbers (default behaviour). - - If `x` is inexact, NaN is replaced by zero or by the user defined value in the - `fill_value` keyword, infinity is replaced by the largest finite floating - point value representable by ``x.dtype``, and -infinity is replaced by the - most negative finite floating point value representable by ``x.dtype``. - - For complex dtypes, the above is applied to each of the real and - imaginary components of `x` separately. - - Parameters - ---------- - x : array | float | complex - Input data. - fill_value : int | float, optional - Value to be used to fill NaN values. If no value is passed - then NaN values will be replaced with 0.0. - xp : array_namespace, optional - The standard-compatible namespace for `x`. Default: infer. - - Returns - ------- - array - `x`, with the non-finite values replaced. - - See Also - -------- - array_api.isnan : Shows which elements are Not a Number (NaN). - - Examples - -------- - >>> import array_api_extra as xpx - >>> import array_api_strict as xp - >>> xpx.nan_to_num(xp.inf) - 1.7976931348623157e+308 - >>> xpx.nan_to_num(-xp.inf) - -1.7976931348623157e+308 - >>> xpx.nan_to_num(xp.nan) - 0.0 - >>> x = xp.asarray([xp.inf, -xp.inf, xp.nan, -128, 128]) - >>> xpx.nan_to_num(x) - array([ 1.79769313e+308, -1.79769313e+308, 0.00000000e+000, # may vary - -1.28000000e+002, 1.28000000e+002]) - >>> y = xp.asarray([complex(xp.inf, xp.nan), xp.nan, complex(xp.nan, xp.inf)]) - array([ 1.79769313e+308, -1.79769313e+308, 0.00000000e+000, # may vary - -1.28000000e+002, 1.28000000e+002]) - >>> xpx.nan_to_num(y) - array([ 1.79769313e+308 +0.00000000e+000j, # may vary - 0.00000000e+000 +0.00000000e+000j, - 0.00000000e+000 +1.79769313e+308j]) - """ - if isinstance(fill_value, complex): - msg = "Complex fill values are not supported." - raise TypeError(msg) - - xp = array_namespace(x) if xp is None else xp - - # for scalars we want to output an array - y = xp.asarray(x) - - if ( - is_cupy_namespace(xp) - or is_jax_namespace(xp) - or is_numpy_namespace(xp) - or is_torch_namespace(xp) - ): - return xp.nan_to_num(y, nan=fill_value) - - return _funcs.nan_to_num(y, fill_value=fill_value, xp=xp) - - -def nunique(x: Array, /, *, xp: ArrayNamespace | None = None) -> Array: - """ - Count the number of unique elements in an array. - - Compatible with JAX and Dask, whose laziness would be otherwise - problematic. - - Parameters - ---------- - x : Array - Input array. - xp : array_namespace, optional - The standard-compatible namespace for `x`. Default: infer. - - Returns - ------- - array: 0-dimensional integer array - The number of unique elements in `x`. It can be lazy. - """ - if xp is None: - xp = array_namespace(x) - - if is_jax_array(x): - # size= is JAX-specific - # https://github.com/data-apis/array-api/issues/883 - _, counts = xp.unique_counts(x, size=size(x)) - return (counts > 0).sum() - - if ( - is_numpy_namespace(xp) - or is_cupy_namespace(xp) - or ( - is_torch_namespace(xp) - and capabilities(xp, device=get_device(x))["data-dependent shapes"] - ) - ): - _, counts = xp.unique_counts(x) - return xp.asarray(size(counts), device=get_device(x)) - - return _funcs.nunique(x, xp=xp) - - -def one_hot( - x: Array, - /, - num_classes: int, - *, - dtype: DType | None = None, - axis: int = -1, - xp: ArrayNamespace | None = None, -) -> Array: - """ - One-hot encode the given indices. - - Each index in the input `x` is encoded as a vector of zeros of length `num_classes` - with the element at the given index set to one. - - Parameters - ---------- - x : array - An array with integral dtype whose values are between `0` and `num_classes - 1`. - num_classes : int - Number of classes in the one-hot dimension. - dtype : DType, optional - The dtype of the return value. Defaults to the default float dtype (usually - float64). - axis : int, optional - Position in the expanded axes where the new axis is placed. Default: -1. - xp : array_namespace, optional - The standard-compatible namespace for `x`. Default: infer. - - Returns - ------- - array - An array having the same shape as `x` except for a new axis at the position - given by `axis` having size `num_classes`. If `axis` is unspecified, it - defaults to -1, which appends a new axis. - - If ``x < 0`` or ``x >= num_classes``, then the result is undefined, may raise - an exception, or may even cause a bad state. `x` is not checked. - - Examples - -------- - >>> import array_api_extra as xpx - >>> import array_api_strict as xp - >>> xpx.one_hot(xp.asarray([1, 2, 0]), 3) - Array([[0., 1., 0.], - [0., 0., 1.], - [1., 0., 0.]], dtype=array_api_strict.float64) - """ - # Validate inputs. - if xp is None: - xp = array_namespace(x) - if not xp.isdtype(x.dtype, "integral"): - msg = "x must have an integral dtype." - raise TypeError(msg) - if dtype is None: - dtype = _funcs.default_dtype(xp, device=get_device(x)) - # Delegate where possible. - if is_jax_namespace(xp): - from jax.nn import one_hot as jax_one_hot - - return jax_one_hot(x, num_classes, dtype=dtype, axis=axis) - if is_torch_namespace(xp): - from torch.nn.functional import one_hot as torch_one_hot - - x = xp.astype(x, xp.int64) # PyTorch only supports int64 here. - try: - out = torch_one_hot(x, num_classes) - except RuntimeError as e: - raise IndexError from e - else: - out = _funcs.one_hot(x, num_classes, xp=xp) - out = xp.astype(out, dtype, copy=False) - if axis != -1: - out = xp.moveaxis(out, -1, axis) - return out - - -def pad( - x: Array, - pad_width: int | tuple[int, int] | Sequence[tuple[int, int]], - mode: Literal["constant"] = "constant", - *, - constant_values: complex = 0, - xp: ArrayNamespace | None = None, -) -> Array: - """ - Pad the input array. - - Parameters - ---------- - x : array - Input array. - pad_width : int or tuple of ints or sequence of pairs of ints - Pad the input array with this many elements from each side. - If a sequence of tuples, ``[(before_0, after_0), ... (before_N, after_N)]``, - each pair applies to the corresponding axis of ``x``. - A single tuple, ``(before, after)``, is equivalent to a list of ``x.ndim`` - copies of this tuple. - mode : str, optional - Only "constant" mode is currently supported, which pads with - the value passed to `constant_values`. - constant_values : python scalar, optional - Use this value to pad the input. Default is zero. - xp : array_namespace, optional - The standard-compatible namespace for `x`. Default: infer. - - Returns - ------- - array - The input array, - padded with ``pad_width`` elements equal to ``constant_values``. - """ - xp = array_namespace(x) if xp is None else xp - - if mode != "constant": - msg = "Only `'constant'` mode is currently supported" - raise NotImplementedError(msg) - - if ( - is_numpy_namespace(xp) - or is_cupy_namespace(xp) - or is_jax_namespace(xp) - or is_pydata_sparse_namespace(xp) - ): - return xp.pad(x, pad_width, mode, constant_values=constant_values) - - if is_torch_namespace(xp): - # normalize `pad_width` on the host rather than through a tensor as done in - # `torch/_numpy`'s implementation (avoids device transfers) - pad_width_seq = normalize_pad_width(pad_width, x.ndim) - # torch.nn.functional.pad counts dimensions from the last one - flat_pad_width = [w for pair in reversed(pad_width_seq) for w in pair] - return xp.nn.functional.pad(x, tuple(flat_pad_width), value=constant_values) - - return _funcs.pad(x, pad_width, constant_values=constant_values, xp=xp) - - -def searchsorted( - x1: Array, - x2: Array, - /, - *, - side: Literal["left", "right"] = "left", - xp: ArrayNamespace | None = None, -) -> Array: - """ - Find indices where elements should be inserted to maintain order. - - Find the indices into a sorted array ``x1`` such that if the elements in ``x2`` - were inserted before the indices, the resulting array would remain sorted. - - The behavior of this function is similar to that of :func:`array_api.searchsorted`, - but it relaxes the requirement that `x1` must be one-dimensional. - This function is vectorized, treating slices along the last axis - as elements and preceding axes as batch (or "loop") dimensions. - - Parameters - ---------- - x1 : Array - Input array. Should have a real-valued data type. Must be sorted in ascending - order along the last axis. - x2 : Array - Array containing search values. Should have a real-valued data type. Must have - the same shape as ``x1`` except along the last axis. - side : {'left', 'right'}, optional - Argument controlling which index is returned if an element of ``x2`` is equal to - one or more elements of ``x1``: ``'left'`` returns the index of the first of - these elements; ``'right'`` returns the next index after the last of these - elements. Default: ``'left'``. - xp : array_namespace, optional - The standard-compatible namespace for the array arguments. Default: infer. - - Returns - ------- - Array: integer array - An array of indices with the same shape as ``x2``. - - Examples - -------- - >>> import array_api_strict as xp - >>> import array_api_extra as xpx - >>> x = xp.asarray([11, 12, 13, 13, 14, 15]) - >>> xpx.searchsorted(x, xp.asarray([10, 11.5, 14.5, 16]), xp=xp) - Array([0, 1, 5, 6], dtype=array_api_strict.int64) - >>> xpx.searchsorted(x, xp.asarray(13), xp=xp) - Array(2, dtype=array_api_strict.int64) - >>> xpx.searchsorted(x, xp.asarray(13), side='right', xp=xp) - Array(4, dtype=array_api_strict.int64) - - `searchsorted` is vectorized along the last axis. - - >>> x1 = xp.asarray([[1., 2., 3., 4.], [5., 6., 7., 8.]]) - >>> x2 = xp.asarray([[1.1, 3.3], [6.6, 8.8]]) - >>> xpx.searchsorted(x1, x2, xp=xp) - Array([[1, 3], - [2, 4]], dtype=array_api_strict.int64) - """ - if xp is None: - xp = array_namespace(x1, x2) - - if side not in {"left", "right"}: - message = "`side` must be either 'left' or 'right'." - raise ValueError(message) - - xp_default_int = _funcs.default_dtype(xp, kind="integral") - x2_0d = x2.ndim == 0 - x1_1d = x1.ndim <= 1 - - if x1_1d or is_torch_namespace(xp): - x2 = xp.reshape(x2, ()) if (x2_0d and x1_1d) else x2 - out = xp.searchsorted(x1, x2, side=side) - return xp.astype(out, xp_default_int, copy=False) - - return _funcs.searchsorted(x1, x2, side=side, xp=xp) - - -def setdiff1d( - x1: Array | complex, - x2: Array | complex, - /, - *, - assume_unique: bool = False, - xp: ArrayNamespace | None = None, -) -> Array: - """ - Find the set difference of two arrays. - - Return the unique values in `x1` that are not in `x2`. - - Parameters - ---------- - x1 : array | int | float | complex | bool - 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, optional - The standard-compatible namespace for `x1` and `x2`. Default: infer. - - Returns - ------- - 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 xp is None: - xp = array_namespace(x1, x2) - - if is_numpy_namespace(xp) or is_cupy_namespace(xp) or is_jax_namespace(xp): - x1, x2 = asarrays(x1, x2, xp=xp) - return xp.setdiff1d(x1, x2, assume_unique=assume_unique) - - return _funcs.setdiff1d(x1, x2, assume_unique=assume_unique, xp=xp) - - -def sinc(x: Array, /, *, xp: ArrayNamespace | None = None) -> Array: - r""" - Return the normalized sinc function. - - The sinc function is equal to :math:`\sin(\pi x)/(\pi x)` for any argument - :math:`x\ne 0`. ``sinc(0)`` takes the limit value 1, making ``sinc`` not - only everywhere continuous but also infinitely differentiable. - - .. note:: - - Note the normalization factor of ``pi`` used in the definition. - This is the most commonly used definition in signal processing. - Use ``sinc(x / xp.pi)`` to obtain the unnormalized sinc function - :math:`\sin(x)/x` that is more common in mathematics. - - Parameters - ---------- - x : array - Array (possibly multi-dimensional) of values for which to calculate - ``sinc(x)``. Must have a real floating point dtype. - xp : array_namespace, optional - The standard-compatible namespace for `x`. Default: infer. - - Returns - ------- - array - ``sinc(x)`` calculated elementwise, which has the same shape as the input. - - Notes - ----- - The name sinc is short for "sine cardinal" or "sinus cardinalis". - - The sinc function is used in various signal processing applications, - including in anti-aliasing, in the construction of a Lanczos resampling - filter, and in interpolation. - - For bandlimited interpolation of discrete-time signals, the ideal - interpolation kernel is proportional to the sinc function. - - References - ---------- - #. 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 - -------- - >>> import array_api_strict as xp - >>> import array_api_extra as xpx - >>> x = xp.linspace(-4, 4, 41) - >>> xpx.sinc(x, xp=xp) - Array([-3.89817183e-17, -4.92362781e-02, - -8.40918587e-02, -8.90384387e-02, - -5.84680802e-02, 3.89817183e-17, - 6.68206631e-02, 1.16434881e-01, - 1.26137788e-01, 8.50444803e-02, - -3.89817183e-17, -1.03943254e-01, - -1.89206682e-01, -2.16236208e-01, - -1.55914881e-01, 3.89817183e-17, - 2.33872321e-01, 5.04551152e-01, - 7.56826729e-01, 9.35489284e-01, - 1.00000000e+00, 9.35489284e-01, - 7.56826729e-01, 5.04551152e-01, - 2.33872321e-01, 3.89817183e-17, - -1.55914881e-01, -2.16236208e-01, - -1.89206682e-01, -1.03943254e-01, - -3.89817183e-17, 8.50444803e-02, - 1.26137788e-01, 1.16434881e-01, - 6.68206631e-02, 3.89817183e-17, - -5.84680802e-02, -8.90384387e-02, - -8.40918587e-02, -4.92362781e-02, - -3.89817183e-17], dtype=array_api_strict.float64) - """ - - if xp is None: - xp = array_namespace(x) - - if not xp.isdtype(x.dtype, "real floating"): - err_msg = "`x` must have a real floating data type." - raise ValueError(err_msg) - - if ( - is_numpy_namespace(xp) - or is_cupy_namespace(xp) - or is_jax_namespace(xp) - or is_torch_namespace(xp) - or is_dask_namespace(xp) - ): - return xp.sinc(x) - - return _funcs.sinc(x, xp=xp) - - -def partition( - a: Array, - kth: int, - /, - axis: int | None = -1, - *, - xp: ArrayNamespace | None = None, -) -> Array: - """ - Return a partitioned copy of an array. - - Creates a copy of the array and partially sorts it in such a way that the value - of the element in k-th position is in the position it would be in a sorted array. - In the output array, all elements smaller than the k-th element are located to - the left of this element and all equal or greater are located to its right. - The ordering of the elements in the two partitions on the either side of - the k-th element in the output array is undefined. - - Parameters - ---------- - a : Array - Input array. - kth : int - Element index to partition by. - axis : int, optional - Axis along which to partition. The default is ``-1`` (the last axis). - If ``None``, the flattened array is used. - xp : array_namespace, optional - The standard-compatible namespace for `x`. Default: infer. - - Returns - ------- - partitioned_array - Array of the same type and shape as `a`. - - Notes - ----- - If `xp` implements ``partition`` or an equivalent function - (e.g. ``topk`` for torch), complexity will likely be O(n). - If not, this function simply calls ``xp.sort`` and complexity is O(n log n). - """ - # Validate inputs. - if xp is None: - xp = array_namespace(a) - if a.ndim < 1: - msg = "`a` must be at least 1-dimensional" - raise TypeError(msg) - if axis is None: - return partition(xp.reshape(a, (-1,)), kth, axis=0, xp=xp) - (size,) = eager_shape(a, axis) - if not (0 <= kth < size): - msg = f"kth(={kth}) out of bounds [0 {size})" - raise ValueError(msg) - - # Delegate where possible. - if is_numpy_namespace(xp) or is_cupy_namespace(xp) or is_jax_namespace(xp): - return xp.partition(a, kth, axis=axis) - - # Use top-k when possible: - if is_torch_namespace(xp): - if not (axis == -1 or axis == a.ndim - 1): - a = xp.transpose(a, axis, -1) - - out = xp.empty_like(a) - ranks = xp.arange(a.shape[-1]).expand_as(a) - - split_value, indices = xp.kthvalue(a, kth + 1, keepdim=True) - del indices # indices won't be used => del ASAP to reduce peak memory usage - - # fill the left-side of the partition - mask_src = a < split_value - n_left = mask_src.sum(dim=-1, keepdim=True) - mask_dest = ranks < n_left - out[mask_dest] = a[mask_src] - - # fill the middle of the partition - mask_src = a == split_value - n_left += mask_src.sum(dim=-1, keepdim=True) - mask_dest ^= ranks < n_left - out[mask_dest] = a[mask_src] - - # fill the right-side of the partition - mask_src = a > split_value - mask_dest = ranks >= n_left - out[mask_dest] = a[mask_src] - - if not (axis == -1 or axis == a.ndim - 1): - out = xp.transpose(out, axis, -1) - return out - - # Note: dask topk/argtopk sort the return values, so it's - # not much more efficient than sorting everything when - # kth is not small compared to x.size - - return _funcs.partition(a, kth, axis=axis, xp=xp) - - -def argpartition( - a: Array, - kth: int, - /, - axis: int | None = -1, - *, - xp: ArrayNamespace | None = None, -) -> Array: - """ - Perform an indirect partition along the given axis. - - It returns an array of indices of the same shape as `a` that - index data along the given axis in partitioned order. - - Parameters - ---------- - a : Array - Input array. - kth : int - Element index to partition by. - axis : int, optional - Axis along which to partition. The default is ``-1`` (the last axis). - If ``None``, the flattened array is used. - xp : array_namespace, optional - The standard-compatible namespace for `x`. Default: infer. - - Returns - ------- - index_array - Array of indices that partition `a` along the specified axis. - - Notes - ----- - If `xp` implements ``argpartition`` or an equivalent function - e.g. ``topk`` for torch), complexity will likely be O(n). - If not, this function simply calls ``xp.argsort`` and complexity is O(n log n). - """ - # Validate inputs. - if xp is None: - xp = array_namespace(a) - if is_pydata_sparse_namespace(xp): - msg = "Not implemented for sparse backend: no argsort" - raise NotImplementedError(msg) - if a.ndim < 1: - msg = "`a` must be at least 1-dimensional" - raise TypeError(msg) - if axis is None: - return argpartition(xp.reshape(a, (-1,)), kth, axis=0, xp=xp) - (size,) = eager_shape(a, axis) - if not (0 <= kth < size): - msg = f"kth(={kth}) out of bounds [0 {size})" - raise ValueError(msg) - - # Delegate where possible. - if is_numpy_namespace(xp) or is_cupy_namespace(xp) or is_jax_namespace(xp): - return xp.argpartition(a, kth, axis=axis) - - # Use top-k when possible: - if is_torch_namespace(xp): - # see `partition` above for commented details of those steps: - if not (axis == -1 or axis == a.ndim - 1): - a = xp.transpose(a, axis, -1) - - ranks = xp.arange(a.shape[-1]).expand_as(a) - out = xp.empty_like(ranks) - - split_value, indices = xp.kthvalue(a, kth + 1, keepdim=True) - del indices # indices won't be used => del ASAP to reduce peak memory usage - - mask_src = a < split_value - n_left = mask_src.sum(dim=-1, keepdim=True) - mask_dest = ranks < n_left - out[mask_dest] = ranks[mask_src] - - mask_src = a == split_value - n_left += mask_src.sum(dim=-1, keepdim=True) - mask_dest ^= ranks < n_left - out[mask_dest] = ranks[mask_src] - - mask_src = a > split_value - mask_dest = ranks >= n_left - out[mask_dest] = ranks[mask_src] - - if not (axis == -1 or axis == a.ndim - 1): - out = xp.transpose(out, axis, -1) - return out - - # Note: dask topk/argtopk sort the return values, so it's - # not much more efficient than sorting everything when - # kth is not small compared to x.size - - return _funcs.argpartition(a, kth, axis=axis, xp=xp) - - -def isin( - a: Array, - b: Array, - /, - *, - assume_unique: bool = False, - invert: bool = False, - kind: str | None = None, - xp: ArrayNamespace | None = None, -) -> Array: - """ - Determine whether each element in `a` is present in `b`. - - This is :func:`array_api.isin`, with additional `assume_unique` - and `kind` parameters. - - Parameters - ---------- - a : array - Input elements. - b : array - The elements against which to test each element of `a`. - assume_unique : bool, optional - If True, the input arrays are both assumed to be unique which can speed - up the calculation. Default: False. - invert : bool, optional - If True, the values in the returned array are inverted. Default: False. - kind : str | None, optional - The algorithm or method to use. This will not affect the final result, - but will affect the speed and memory use. - For NumPy the options are {None, "sort", "table"}. - For Jax the mapped parameter is instead `method` and the options are - {"compare_all", "binary_search", "sort", and "auto" (default)} - For CuPy, Dask, Torch and the default case this parameter is not present and - thus ignored. Default: None. - xp : array_namespace, optional - The standard-compatible namespace for `a` and `b`. Default: infer. - - Returns - ------- - array - An array having the same shape as that of `a` that is True for elements - that are in `b` and False otherwise. - """ - if xp is None: - xp = array_namespace(a, b) - - if is_numpy_namespace(xp): - return xp.isin(a, b, assume_unique=assume_unique, invert=invert, kind=kind) - if is_jax_namespace(xp): - if kind is None: - kind = "auto" - return xp.isin(a, b, assume_unique=assume_unique, invert=invert, method=kind) - if is_cupy_namespace(xp) or is_torch_namespace(xp) or is_dask_namespace(xp): - return xp.isin(a, b, assume_unique=assume_unique, invert=invert) - - return _funcs.isin(a, b, assume_unique=assume_unique, invert=invert, xp=xp) - - -def union1d(a: Array, b: Array, /, *, xp: ArrayNamespace | None = None) -> Array: - """ - Find the union of two arrays. - - Return the unique, sorted array of values that are in either of the two - input arrays. - - Parameters - ---------- - a, b : Array - Input arrays. They are flattened internally if they are not already 1D. - - xp : array_namespace, optional - The standard-compatible namespace for `a` and `b`. Default: infer. - - Returns - ------- - Array - Unique, sorted union of the input arrays. - - See Also - -------- - jax.numpy.union1d : Corresponding function in JAX. - - Notes - ----- - This function is not compatible with `jax.jit`. - See the docstring of the corresponding JAX function for more information. - """ - if xp is None: - xp = array_namespace(a, b) - - if ( - is_numpy_namespace(xp) - or is_cupy_namespace(xp) - or is_dask_namespace(xp) - or is_jax_namespace(xp) - ): - return xp.union1d(a, b) - - return _funcs.union1d(a, b, xp=xp) - - -def tril_indices( - n: int, - /, - *, - offset: int = 0, - m: int | None = None, - device: Device | None = None, - xp: ArrayNamespace, -) -> tuple[Array, Array]: - """ - Return the indices of the lower triangle of an ``(n, m)`` array. - - Equivalent to :func:`numpy.tril_indices` with parameter ``k`` renamed to - ``offset`` to match :func:`array_api.linalg.diagonal`'s naming. - - Parameters - ---------- - n : int - The row dimension of the array. - offset : int, optional - Diagonal offset; ``0`` (default) is the main diagonal. Corresponds - to ``k`` in :func:`numpy.tril_indices`. - m : int, optional - The column dimension. If ``None`` (default), assumed equal to `n`. - device : Device, optional - The device on which to place the returned arrays. Default: current device. - xp : array_namespace - The standard-compatible namespace to create the indices in. - - Returns - ------- - tuple of array - Row and column indices ``(rows, cols)`` of the lower triangle of - the ``(n, m)`` matrix, shifted by `offset`. - - Notes - ----- - The generic fallback uses :func:`array_api.nonzero`, so namespaces without - ``nonzero`` are not supported on that path. - - Examples - -------- - >>> import array_api_strict as xp - >>> import array_api_extra as xpx - >>> rows, cols = xpx.tril_indices(3, xp=xp) - >>> rows - Array([0, 1, 1, 2, 2, 2], dtype=array_api_strict.int64) - >>> cols - Array([0, 0, 1, 0, 1, 2], dtype=array_api_strict.int64) - """ - if n < 0: - msg = f"`n` must be non-negative, got {n}" - raise ValueError(msg) - if m is not None and m < 0: - msg = f"`m` must be non-negative, got {m}" - raise ValueError(msg) - if device is None and ( - is_numpy_namespace(xp) - or is_cupy_namespace(xp) - or is_jax_namespace(xp) - or is_dask_namespace(xp) - ): - return xp.tril_indices(n, k=offset, m=m) - if is_torch_namespace(xp): - # `torch.tril_indices` returns a 2xN tensor, not a tuple, and - # takes (row, col) rather than (n, *, m=None). - cols = n if m is None else m - idx = xp.tril_indices(n, cols, offset=offset, device=device) - return (idx[0], idx[1]) - return _funcs.tril_indices(n, offset=offset, m=m, device=device, xp=xp) - - -def triu_indices( - n: int, - /, - *, - offset: int = 0, - m: int | None = None, - device: Device | None = None, - xp: ArrayNamespace, -) -> tuple[Array, Array]: - """ - Return the indices of the upper triangle of an ``(n, m)`` array. - - Equivalent to :func:`numpy.triu_indices` with parameter ``k`` renamed to - ``offset`` to match :func:`array_api.linalg.diagonal`'s naming. - - Parameters - ---------- - n : int - The row dimension of the array. - offset : int, optional - Diagonal offset; ``0`` (default) is the main diagonal. Corresponds - to ``k`` in :func:`numpy.triu_indices`. - m : int, optional - The column dimension. If ``None`` (default), assumed equal to `n`. - device : Device, optional - The device on which to place the returned arrays. Default: current device. - xp : array_namespace - The standard-compatible namespace to create the indices in. - - Returns - ------- - tuple of array - Row and column indices ``(rows, cols)`` of the upper triangle of - the ``(n, m)`` matrix, shifted by `offset`. - - Notes - ----- - The generic fallback uses :func:`array_api.nonzero`, so namespaces without - ``nonzero`` are not supported on that path. - - Examples - -------- - >>> import array_api_strict as xp - >>> import array_api_extra as xpx - >>> rows, cols = xpx.triu_indices(3, xp=xp) - >>> rows - Array([0, 0, 0, 1, 1, 2], dtype=array_api_strict.int64) - >>> cols - Array([0, 1, 2, 1, 2, 2], dtype=array_api_strict.int64) - """ - if n < 0: - msg = f"`n` must be non-negative, got {n}" - raise ValueError(msg) - if m is not None and m < 0: - msg = f"`m` must be non-negative, got {m}" - raise ValueError(msg) - if device is None and ( - is_numpy_namespace(xp) - or is_cupy_namespace(xp) - or is_jax_namespace(xp) - or is_dask_namespace(xp) - ): - return xp.triu_indices(n, k=offset, m=m) - if is_torch_namespace(xp): - cols = n if m is None else m - idx = xp.triu_indices(n, cols, offset=offset, device=device) - return (idx[0], idx[1]) - return _funcs.triu_indices(n, offset=offset, m=m, device=device, xp=xp) - - -def unravel_index( - indices: Array, - shape: tuple[int, ...], - /, - *, - xp: ArrayNamespace | None = None, -) -> tuple[Array, ...]: - """ - Convert a flat index or array of flat indices into a tuple of coordinate arrays. - - Parameters - ---------- - indices : array - An integer array whose elements are indices into the flattened version - of an array of dimensions `shape`. - - shape : tuple of ints - The shape to use for unraveling `indices`. - - xp : array_namespace, optional - The standard-compatible namespace for `indices`. Default: infer. - - Returns - ------- - tuple of array - A tuple of unraveled indices. Each array in the tuple has the same shape - as the `indices` array. - - Examples - -------- - >>> import array_api_extra as xpx - >>> import array_api_strict as xp - >>> xs, ys = xpx.unravel_index(xp.asarray([1, 2, 4, 5, 6, 8]), (4, 3)) - >>> xs, ys - ( - Array([0, 0, 1, 1, 2, 2], dtype=array_api_strict.int64), - Array([1, 2, 1, 2, 0, 2], dtype=array_api_strict.int64), - ) - >>> [(int(x), int(y)) for x, y in zip(xs, ys)] - [(0, 1), (0, 2), (1, 1), (1, 2), (2, 0), (2, 2)] - >>> xs, ys = xpx.unravel_index(xp.arange(6), (2, 2)) - >>> [(int(x), int(y)) for x, y in zip(xs, ys)] - [(0, 0), (0, 1), (1, 0), (1, 1), (0, 0), (0, 1)] - """ - if xp is None: - xp = array_namespace(indices) - - if ( - is_numpy_namespace(xp) - or is_cupy_namespace(xp) - or is_dask_namespace(xp) - or is_jax_namespace(xp) - or is_torch_namespace(xp) - ): - return xp.unravel_index(indices, shape) - - return _funcs.unravel_index(indices, shape) - - -def nanmin( - a: Array, - /, - *, - axis: int | tuple[int, ...] | None = None, - xp: ArrayNamespace | None = None, -) -> Array: - """ - Return the minimum of the array elements along a given axis, ignoring NaNs. - - Parameters - ---------- - a : Array - Input array. - axis : int or tuple of ints or None, optional - Axis or axes along which the minimum is computed. The default is to compute - the minimum of the flattened array. - xp : array_namespace, optional - The standard-compatible namespace for `a`. Default: infer. - - Returns - ------- - array - An array of minimum values along the given axis, ignoring NaNs. - - Examples - -------- - >>> import array_api_extra as xpx - >>> import array_api_strict as xp - >>> a = xp.asarray([[5, 3, xp.nan, 1], [4, xp.nan, 2, xp.nan]]) - >>> xpx.nanmin(a) - Array(1., dtype=array_api_strict.float64) - >>> xpx.nanmin(a, axis=0) - Array([4., 3., 2., 1.], dtype=array_api_strict.float64) - >>> xpx.nanmin(a, axis=1) - Array([1., 2.], dtype=array_api_strict.float64) - """ - if xp is None: - xp = array_namespace(a) - - if ( - is_numpy_namespace(xp) - or is_cupy_namespace(xp) - or is_dask_namespace(xp) - or is_jax_namespace(xp) - ): - return xp.nanmin(a, axis=axis) - - return _funcs.nanmin(a, axis=axis, xp=xp) - - -def nanmax( - a: Array, - /, - *, - axis: int | tuple[int, ...] | None = None, - xp: ArrayNamespace | None = None, -) -> Array: - """ - Return the maximum of the array elements along a given axis, ignoring NaNs. - - Parameters - ---------- - a : Array - Input array. - axis : int or tuple of ints or None, optional - Axis or axes along which the maximum is computed. The default is to compute - the maximum of the flattened array. - xp : array_namespace, optional - The standard-compatible namespace for `a`. Default: infer. - - Returns - ------- - array - An array of maximum values along the given axis, ignoring NaNs. - - Examples - -------- - >>> import array_api_extra as xpx - >>> import array_api_strict as xp - >>> a = xp.asarray([[5, 3, xp.nan, 6], [4, xp.nan, 2, xp.nan]]) - >>> xpx.nanmax(a) - Array(6., dtype=array_api_strict.float64) - >>> xpx.nanmax(a, axis=0) - Array([5., 3., 2., 6.], dtype=array_api_strict.float64) - >>> xpx.nanmax(a, axis=1) - Array([6., 4.], dtype=array_api_strict.float64) - """ - if xp is None: - xp = array_namespace(a) - - if ( - is_numpy_namespace(xp) - or is_cupy_namespace(xp) - or is_dask_namespace(xp) - or is_jax_namespace(xp) - ): - return xp.nanmax(a, axis=axis) - - return _funcs.nanmax(a, axis=axis, xp=xp) - - -def nansum( - a: Array, - /, - *, - axis: int | tuple[int, ...] | None = None, - xp: ArrayNamespace | None = None, -) -> Array: - """ - Return the sum of the array elements along a given axis, ignoring NaNs. - - Parameters - ---------- - a : Array - Input array. - axis : int or tuple of ints or None, optional - Axis or axes along which the sum is computed. The default is to compute - the sum of the flattened array. - xp : array_namespace, optional - The standard-compatible namespace for `a`. Default: infer. - - Returns - ------- - array - An array of sum values along the given axis, ignoring NaNs. - - Examples - -------- - >>> import array_api_extra as xpx - >>> import array_api_strict as xp - >>> a = xp.asarray([[5, 3, xp.nan, 1], [4, xp.nan, 2, xp.nan]]) - >>> xpx.nansum(a) - Array(15., dtype=array_api_strict.float64) - >>> xpx.nansum(a, axis=0) - Array([9., 3., 2., 1.], dtype=array_api_strict.float64) - >>> xpx.nansum(a, axis=1) - Array([9., 6.], dtype=array_api_strict.float64) - """ - if xp is None: - xp = array_namespace(a) - - if ( - is_numpy_namespace(xp) - or is_cupy_namespace(xp) - or is_dask_namespace(xp) - or is_jax_namespace(xp) - or is_torch_namespace(xp) - ): - return xp.nansum(a, axis=axis) - - return _funcs.nansum(a, axis=axis, xp=xp) diff --git a/src/array_api_extra/_elementwise.py b/src/array_api_extra/_elementwise.py new file mode 100644 index 00000000..ca1f9dd8 --- /dev/null +++ b/src/array_api_extra/_elementwise.py @@ -0,0 +1,373 @@ +"""Delegation layer for element-wise functions.""" + +from ._agnostic import _elementwise, _inspection +from ._lib._compat import ( + array_namespace, + is_cupy_namespace, + is_dask_namespace, + is_jax_namespace, + is_numpy_namespace, + is_torch_namespace, +) +from ._lib._compat import device as get_device +from ._lib._helpers import asarrays +from ._lib._typing import Array, ArrayNamespace + +__all__ = ["deg2rad", "isclose", "nan_to_num", "rad2deg", "sinc"] + + +def deg2rad(x: Array, /, *, xp: ArrayNamespace | None = None) -> Array: + """ + Convert angles from degrees to radians. + + Parameters + ---------- + x : array + Input array in degrees. Must have an integral or floating-point dtype. + xp : array_namespace, optional + The standard-compatible namespace for `x`. Default: infer. + + Returns + ------- + array + The corresponding angles in radians. Integral inputs are converted to the + default floating-point dtype. + + Examples + -------- + >>> import array_api_strict as xp + >>> import array_api_extra as xpx + >>> xpx.deg2rad(xp.asarray([0, 90, 180]), xp=xp) + Array([0. , 1.57079633, 3.14159265], dtype=array_api_strict.float64) + """ + if xp is None: + xp = array_namespace(x) + if xp.isdtype(x.dtype, "integral"): + x = xp.astype(x, _inspection.default_dtype(xp, device=get_device(x))) + elif not xp.isdtype(x.dtype, ("real floating", "complex floating")): + msg = "`x` must have an integral, real floating, or complex floating dtype." + raise TypeError(msg) + + if is_jax_namespace(xp) or ( + not xp.isdtype(x.dtype, "complex floating") + and ( + is_numpy_namespace(xp) + or is_cupy_namespace(xp) + or is_torch_namespace(xp) + or is_dask_namespace(xp) + ) + ): + return xp.deg2rad(x) + + return _elementwise.deg2rad(x, xp=xp) + + +def rad2deg(x: Array, /, *, xp: ArrayNamespace | None = None) -> Array: + """ + Convert angles from radians to degrees. + + Parameters + ---------- + x : array + Input array in radians. Must have an integral or floating-point dtype. + xp : array_namespace, optional + The standard-compatible namespace for `x`. Default: infer. + + Returns + ------- + array + The corresponding angles in degrees. Integral inputs are converted to the + default floating-point dtype. + + Examples + -------- + >>> import array_api_strict as xp + >>> import array_api_extra as xpx + >>> xpx.rad2deg(xp.asarray([0.0, xp.pi / 2, xp.pi]), xp=xp) + Array([ 0., 90., 180.], dtype=array_api_strict.float64) + """ + if xp is None: + xp = array_namespace(x) + if xp.isdtype(x.dtype, "integral"): + x = xp.astype(x, _inspection.default_dtype(xp, device=get_device(x))) + elif not xp.isdtype(x.dtype, ("real floating", "complex floating")): + msg = "`x` must have an integral, real floating, or complex floating dtype." + raise TypeError(msg) + + if is_jax_namespace(xp) or ( + not xp.isdtype(x.dtype, "complex floating") + and ( + is_numpy_namespace(xp) + or is_cupy_namespace(xp) + or is_torch_namespace(xp) + or is_dask_namespace(xp) + ) + ): + return xp.rad2deg(x) + + return _elementwise.rad2deg(x, xp=xp) + + +def isclose( + a: Array | complex, + b: Array | complex, + *, + rtol: float = 1e-05, + atol: float = 1e-08, + equal_nan: bool = False, + xp: ArrayNamespace | None = None, +) -> Array: + """ + Return a boolean array where two arrays are element-wise equal within a tolerance. + + The tolerance values are positive, typically very small numbers. The relative + difference ``(rtol * abs(b))`` and the absolute difference `atol` are added together + to compare against the absolute difference between `a` and `b`. + + NaNs are treated as equal if they are in the same place and if ``equal_nan=True``. + Infs are treated as equal if they are in the same place and of the same sign in both + arrays. + + Parameters + ---------- + a, b : Array | int | float | complex | bool + Input objects to compare. At least one must be an array. + rtol : array_like, optional + The relative tolerance parameter (see Notes). + atol : array_like, optional + The absolute tolerance parameter (see Notes). + equal_nan : bool, optional + Whether to compare NaN's as equal. If True, NaN's in `a` will be considered + equal to NaN's in `b` in the output array. + xp : array_namespace, optional + The standard-compatible namespace for `a` and `b`. Default: infer. + + Returns + ------- + Array + A boolean array of shape broadcasted from `a` and `b`, containing ``True`` where + `a` is close to `b`, and ``False`` otherwise. + + Warnings + -------- + The default `atol` is not appropriate for comparing numbers with magnitudes much + smaller than one (see notes). + + See Also + -------- + math.isclose : Similar function in stdlib for Python scalars. + + Notes + ----- + For finite values, `isclose` uses the following equation to test whether two + floating point values are equivalent:: + + absolute(a - b) <= (atol + rtol * absolute(b)) + + Unlike the built-in `math.isclose`, + the above equation is not symmetric in `a` and `b`, + so that ``isclose(a, b)`` might be different from ``isclose(b, a)`` in some rare + cases. + + The default value of `atol` is not appropriate when the reference value `b` has + magnitude smaller than one. For example, it is unlikely that ``a = 1e-9`` and + ``b = 2e-9`` should be considered "close", yet ``isclose(1e-9, 2e-9)`` is ``True`` + with default settings. Be sure to select `atol` for the use case at hand, especially + for defining the threshold below which a non-zero value in `a` will be considered + "close" to a very small or zero value in `b`. + + The comparison of `a` and `b` uses standard broadcasting, which means that `a` and + `b` need not have the same shape in order for ``isclose(a, b)`` to evaluate to + ``True``. + + `isclose` is not defined for non-numeric data types. + ``bool`` is considered a numeric data-type for this purpose. + """ + xp = array_namespace(a, b) if xp is None else xp + + if ( + is_numpy_namespace(xp) + or is_cupy_namespace(xp) + or is_dask_namespace(xp) + or is_jax_namespace(xp) + ): + return xp.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) + + if is_torch_namespace(xp): + a, b = asarrays(a, b, xp=xp) # Array API 2024.12 support + return xp.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan) + + return _elementwise.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan, xp=xp) + + +def nan_to_num( + x: Array | complex, + /, + *, + fill_value: float = 0.0, + xp: ArrayNamespace | None = None, +) -> Array: + """ + Replace NaN with zero and infinity with large finite numbers (default behaviour). + + If `x` is inexact, NaN is replaced by zero or by the user defined value in the + `fill_value` keyword, infinity is replaced by the largest finite floating + point value representable by ``x.dtype``, and -infinity is replaced by the + most negative finite floating point value representable by ``x.dtype``. + + For complex dtypes, the above is applied to each of the real and + imaginary components of `x` separately. + + Parameters + ---------- + x : array | float | complex + Input data. + fill_value : int | float, optional + Value to be used to fill NaN values. If no value is passed + then NaN values will be replaced with 0.0. + xp : array_namespace, optional + The standard-compatible namespace for `x`. Default: infer. + + Returns + ------- + array + `x`, with the non-finite values replaced. + + See Also + -------- + array_api.isnan : Shows which elements are Not a Number (NaN). + + Examples + -------- + >>> import array_api_extra as xpx + >>> import array_api_strict as xp + >>> xpx.nan_to_num(xp.inf) + 1.7976931348623157e+308 + >>> xpx.nan_to_num(-xp.inf) + -1.7976931348623157e+308 + >>> xpx.nan_to_num(xp.nan) + 0.0 + >>> x = xp.asarray([xp.inf, -xp.inf, xp.nan, -128, 128]) + >>> xpx.nan_to_num(x) + array([ 1.79769313e+308, -1.79769313e+308, 0.00000000e+000, # may vary + -1.28000000e+002, 1.28000000e+002]) + >>> y = xp.asarray([complex(xp.inf, xp.nan), xp.nan, complex(xp.nan, xp.inf)]) + array([ 1.79769313e+308, -1.79769313e+308, 0.00000000e+000, # may vary + -1.28000000e+002, 1.28000000e+002]) + >>> xpx.nan_to_num(y) + array([ 1.79769313e+308 +0.00000000e+000j, # may vary + 0.00000000e+000 +0.00000000e+000j, + 0.00000000e+000 +1.79769313e+308j]) + """ + if isinstance(fill_value, complex): + msg = "Complex fill values are not supported." + raise TypeError(msg) + + xp = array_namespace(x) if xp is None else xp + + # for scalars we want to output an array + y = xp.asarray(x) + + if ( + is_cupy_namespace(xp) + or is_jax_namespace(xp) + or is_numpy_namespace(xp) + or is_torch_namespace(xp) + ): + return xp.nan_to_num(y, nan=fill_value) + + return _elementwise.nan_to_num(y, fill_value=fill_value, xp=xp) + + +def sinc(x: Array, /, *, xp: ArrayNamespace | None = None) -> Array: + r""" + Return the normalized sinc function. + + The sinc function is equal to :math:`\sin(\pi x)/(\pi x)` for any argument + :math:`x\ne 0`. ``sinc(0)`` takes the limit value 1, making ``sinc`` not + only everywhere continuous but also infinitely differentiable. + + .. note:: + + Note the normalization factor of ``pi`` used in the definition. + This is the most commonly used definition in signal processing. + Use ``sinc(x / xp.pi)`` to obtain the unnormalized sinc function + :math:`\sin(x)/x` that is more common in mathematics. + + Parameters + ---------- + x : array + Array (possibly multi-dimensional) of values for which to calculate + ``sinc(x)``. Must have a real floating point dtype. + xp : array_namespace, optional + The standard-compatible namespace for `x`. Default: infer. + + Returns + ------- + array + ``sinc(x)`` calculated elementwise, which has the same shape as the input. + + Notes + ----- + The name sinc is short for "sine cardinal" or "sinus cardinalis". + + The sinc function is used in various signal processing applications, + including in anti-aliasing, in the construction of a Lanczos resampling + filter, and in interpolation. + + For bandlimited interpolation of discrete-time signals, the ideal + interpolation kernel is proportional to the sinc function. + + References + ---------- + #. 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 + -------- + >>> import array_api_strict as xp + >>> import array_api_extra as xpx + >>> x = xp.linspace(-4, 4, 41) + >>> xpx.sinc(x, xp=xp) + Array([-3.89817183e-17, -4.92362781e-02, + -8.40918587e-02, -8.90384387e-02, + -5.84680802e-02, 3.89817183e-17, + 6.68206631e-02, 1.16434881e-01, + 1.26137788e-01, 8.50444803e-02, + -3.89817183e-17, -1.03943254e-01, + -1.89206682e-01, -2.16236208e-01, + -1.55914881e-01, 3.89817183e-17, + 2.33872321e-01, 5.04551152e-01, + 7.56826729e-01, 9.35489284e-01, + 1.00000000e+00, 9.35489284e-01, + 7.56826729e-01, 5.04551152e-01, + 2.33872321e-01, 3.89817183e-17, + -1.55914881e-01, -2.16236208e-01, + -1.89206682e-01, -1.03943254e-01, + -3.89817183e-17, 8.50444803e-02, + 1.26137788e-01, 1.16434881e-01, + 6.68206631e-02, 3.89817183e-17, + -5.84680802e-02, -8.90384387e-02, + -8.40918587e-02, -4.92362781e-02, + -3.89817183e-17], dtype=array_api_strict.float64) + """ + + if xp is None: + xp = array_namespace(x) + + if not xp.isdtype(x.dtype, "real floating"): + err_msg = "`x` must have a real floating data type." + raise ValueError(err_msg) + + if ( + is_numpy_namespace(xp) + or is_cupy_namespace(xp) + or is_jax_namespace(xp) + or is_torch_namespace(xp) + or is_dask_namespace(xp) + ): + return xp.sinc(x) + + return _elementwise.sinc(x, xp=xp) diff --git a/src/array_api_extra/_indexing.py b/src/array_api_extra/_indexing.py new file mode 100644 index 00000000..1af41304 --- /dev/null +++ b/src/array_api_extra/_indexing.py @@ -0,0 +1,264 @@ +"""Delegation layer for indexing functions.""" + +from ._agnostic import _indexing +from ._lib._compat import ( + array_namespace, + is_cupy_namespace, + is_dask_namespace, + is_jax_namespace, + is_numpy_namespace, + is_torch_namespace, +) +from ._lib._typing import Array, ArrayNamespace, Device + +__all__ = ["diag_indices", "tril_indices", "triu_indices", "unravel_index"] + + +def diag_indices( + n: int, /, *, ndim: int = 2, device: Device | None = None, xp: ArrayNamespace +) -> tuple[Array, ...]: + """ + Return the indices to access the main diagonal of an array. + + Equivalent to :func:`numpy.diag_indices`. + + Parameters + ---------- + n : int + The size of each dimension of the (hyper-)cube ``(n, n, ..., n)`` + that the returned indices index into. + ndim : int, optional + The number of dimensions. Default: ``2``. + device : Device, optional + The device on which to place the returned arrays. Default: current device. + xp : array_namespace + The standard-compatible namespace to create the indices in. + + Returns + ------- + tuple of array + 1-D integer arrays of length ``n`` that together index + the main diagonal of an array of shape ``(n,) * ndim``. + + Examples + -------- + >>> import array_api_strict as xp + >>> import array_api_extra as xpx + >>> rows, cols = xpx.diag_indices(3, xp=xp) + >>> rows + Array([0, 1, 2], dtype=array_api_strict.int64) + >>> cols + Array([0, 1, 2], dtype=array_api_strict.int64) + """ + if n < 0: + msg = f"`n` must be non-negative, got {n}" + raise ValueError(msg) + if ndim < 1: + msg = f"`ndim` must be >= 1, got {ndim}" + raise ValueError(msg) + if device is None and ( + is_numpy_namespace(xp) or is_cupy_namespace(xp) or is_jax_namespace(xp) + ): + return xp.diag_indices(n, ndim=ndim) + return _indexing.diag_indices(n, ndim=ndim, device=device, xp=xp) + + +def tril_indices( + n: int, + /, + *, + offset: int = 0, + m: int | None = None, + device: Device | None = None, + xp: ArrayNamespace, +) -> tuple[Array, Array]: + """ + Return the indices of the lower triangle of an ``(n, m)`` array. + + Equivalent to :func:`numpy.tril_indices` with parameter ``k`` renamed to + ``offset`` to match :func:`array_api.linalg.diagonal`'s naming. + + Parameters + ---------- + n : int + The row dimension of the array. + offset : int, optional + Diagonal offset; ``0`` (default) is the main diagonal. Corresponds + to ``k`` in :func:`numpy.tril_indices`. + m : int, optional + The column dimension. If ``None`` (default), assumed equal to `n`. + device : Device, optional + The device on which to place the returned arrays. Default: current device. + xp : array_namespace + The standard-compatible namespace to create the indices in. + + Returns + ------- + tuple of array + Row and column indices ``(rows, cols)`` of the lower triangle of + the ``(n, m)`` matrix, shifted by `offset`. + + Notes + ----- + The generic fallback uses :func:`array_api.nonzero`, so namespaces without + ``nonzero`` are not supported on that path. + + Examples + -------- + >>> import array_api_strict as xp + >>> import array_api_extra as xpx + >>> rows, cols = xpx.tril_indices(3, xp=xp) + >>> rows + Array([0, 1, 1, 2, 2, 2], dtype=array_api_strict.int64) + >>> cols + Array([0, 0, 1, 0, 1, 2], dtype=array_api_strict.int64) + """ + if n < 0: + msg = f"`n` must be non-negative, got {n}" + raise ValueError(msg) + if m is not None and m < 0: + msg = f"`m` must be non-negative, got {m}" + raise ValueError(msg) + if device is None and ( + is_numpy_namespace(xp) + or is_cupy_namespace(xp) + or is_jax_namespace(xp) + or is_dask_namespace(xp) + ): + return xp.tril_indices(n, k=offset, m=m) + if is_torch_namespace(xp): + # `torch.tril_indices` returns a 2xN tensor, not a tuple, and + # takes (row, col) rather than (n, *, m=None). + cols = n if m is None else m + idx = xp.tril_indices(n, cols, offset=offset, device=device) + return (idx[0], idx[1]) + return _indexing.tril_indices(n, offset=offset, m=m, device=device, xp=xp) + + +def triu_indices( + n: int, + /, + *, + offset: int = 0, + m: int | None = None, + device: Device | None = None, + xp: ArrayNamespace, +) -> tuple[Array, Array]: + """ + Return the indices of the upper triangle of an ``(n, m)`` array. + + Equivalent to :func:`numpy.triu_indices` with parameter ``k`` renamed to + ``offset`` to match :func:`array_api.linalg.diagonal`'s naming. + + Parameters + ---------- + n : int + The row dimension of the array. + offset : int, optional + Diagonal offset; ``0`` (default) is the main diagonal. Corresponds + to ``k`` in :func:`numpy.triu_indices`. + m : int, optional + The column dimension. If ``None`` (default), assumed equal to `n`. + device : Device, optional + The device on which to place the returned arrays. Default: current device. + xp : array_namespace + The standard-compatible namespace to create the indices in. + + Returns + ------- + tuple of array + Row and column indices ``(rows, cols)`` of the upper triangle of + the ``(n, m)`` matrix, shifted by `offset`. + + Notes + ----- + The generic fallback uses :func:`array_api.nonzero`, so namespaces without + ``nonzero`` are not supported on that path. + + Examples + -------- + >>> import array_api_strict as xp + >>> import array_api_extra as xpx + >>> rows, cols = xpx.triu_indices(3, xp=xp) + >>> rows + Array([0, 0, 0, 1, 1, 2], dtype=array_api_strict.int64) + >>> cols + Array([0, 1, 2, 1, 2, 2], dtype=array_api_strict.int64) + """ + if n < 0: + msg = f"`n` must be non-negative, got {n}" + raise ValueError(msg) + if m is not None and m < 0: + msg = f"`m` must be non-negative, got {m}" + raise ValueError(msg) + if device is None and ( + is_numpy_namespace(xp) + or is_cupy_namespace(xp) + or is_jax_namespace(xp) + or is_dask_namespace(xp) + ): + return xp.triu_indices(n, k=offset, m=m) + if is_torch_namespace(xp): + cols = n if m is None else m + idx = xp.triu_indices(n, cols, offset=offset, device=device) + return (idx[0], idx[1]) + return _indexing.triu_indices(n, offset=offset, m=m, device=device, xp=xp) + + +def unravel_index( + indices: Array, + shape: tuple[int, ...], + /, + *, + xp: ArrayNamespace | None = None, +) -> tuple[Array, ...]: + """ + Convert a flat index or array of flat indices into a tuple of coordinate arrays. + + Parameters + ---------- + indices : array + An integer array whose elements are indices into the flattened version + of an array of dimensions `shape`. + + shape : tuple of ints + The shape to use for unraveling `indices`. + + xp : array_namespace, optional + The standard-compatible namespace for `indices`. Default: infer. + + Returns + ------- + tuple of array + A tuple of unraveled indices. Each array in the tuple has the same shape + as the `indices` array. + + Examples + -------- + >>> import array_api_extra as xpx + >>> import array_api_strict as xp + >>> xs, ys = xpx.unravel_index(xp.asarray([1, 2, 4, 5, 6, 8]), (4, 3)) + >>> xs, ys + ( + Array([0, 0, 1, 1, 2, 2], dtype=array_api_strict.int64), + Array([1, 2, 1, 2, 0, 2], dtype=array_api_strict.int64), + ) + >>> [(int(x), int(y)) for x, y in zip(xs, ys)] + [(0, 1), (0, 2), (1, 1), (1, 2), (2, 0), (2, 2)] + >>> xs, ys = xpx.unravel_index(xp.arange(6), (2, 2)) + >>> [(int(x), int(y)) for x, y in zip(xs, ys)] + [(0, 0), (0, 1), (1, 0), (1, 1), (0, 0), (0, 1)] + """ + if xp is None: + xp = array_namespace(indices) + + if ( + is_numpy_namespace(xp) + or is_cupy_namespace(xp) + or is_dask_namespace(xp) + or is_jax_namespace(xp) + or is_torch_namespace(xp) + ): + return xp.unravel_index(indices, shape) + + return _indexing.unravel_index(indices, shape) diff --git a/src/array_api_extra/_lib/_lazy.py b/src/array_api_extra/_lazy.py similarity index 98% rename from src/array_api_extra/_lib/_lazy.py rename to src/array_api_extra/_lazy.py index d55e9c21..50f98134 100644 --- a/src/array_api_extra/_lib/_lazy.py +++ b/src/array_api_extra/_lazy.py @@ -1,4 +1,4 @@ -"""Public API Functions.""" +"""Tools for lazy backends.""" from __future__ import annotations @@ -7,15 +7,15 @@ from functools import partial, wraps from typing import TYPE_CHECKING, Any, ParamSpec, TypeAlias, cast, overload -from ._funcs import broadcast_shapes -from ._utils import _compat -from ._utils._compat import ( +from ._agnostic._manipulation import broadcast_shapes +from ._lib import _compat +from ._lib._compat import ( array_namespace, is_dask_namespace, is_jax_namespace, ) -from ._utils._helpers import is_python_scalar -from ._utils._typing import Array, ArrayNamespace, DType +from ._lib._helpers import is_python_scalar +from ._lib._typing import Array, ArrayNamespace, DType if TYPE_CHECKING: # pragma: no cover import numpy as np @@ -28,6 +28,8 @@ P = ParamSpec("P") +__all__ = ["lazy_apply"] + @overload # pyrefly: ignore[invalid-param-spec] def lazy_apply( # type: ignore[valid-type] diff --git a/src/array_api_extra/_lib/_utils/_compat.py b/src/array_api_extra/_lib/_compat.py similarity index 97% rename from src/array_api_extra/_lib/_utils/_compat.py rename to src/array_api_extra/_lib/_compat.py index 82ce76b8..98f213ba 100644 --- a/src/array_api_extra/_lib/_utils/_compat.py +++ b/src/array_api_extra/_lib/_compat.py @@ -4,7 +4,7 @@ # pylint: disable=duplicate-code try: - from ...._array_api_compat_vendor import ( + from ..._array_api_compat_vendor import ( array_namespace, device, is_array_api_obj, diff --git a/src/array_api_extra/_lib/_utils/_compat.pyi b/src/array_api_extra/_lib/_compat.pyi similarity index 100% rename from src/array_api_extra/_lib/_utils/_compat.pyi rename to src/array_api_extra/_lib/_compat.pyi diff --git a/src/array_api_extra/_lib/_funcs.py b/src/array_api_extra/_lib/_funcs.py deleted file mode 100644 index 4629461e..00000000 --- a/src/array_api_extra/_lib/_funcs.py +++ /dev/null @@ -1,943 +0,0 @@ -"""Array-agnostic implementations for the public API.""" - -import math -import warnings -from collections.abc import Callable, Sequence -from types import NoneType -from typing import Literal, cast, overload - -from ._at import at -from ._utils import _compat, _helpers -from ._utils._compat import ( - array_namespace, - is_dask_namespace, -) -from ._utils._helpers import ( - asarrays, - capabilities, - eager_shape, - meta_namespace, - ndindex, - normalize_pad_width, -) -from ._utils._typing import Array, ArrayNamespace, Device, DType - -__all__ = [ - "angle", - "apply_where", - "argpartition", - "atleast_nd", - "broadcast_shapes", - "cov", - "create_diagonal", - "default_dtype", - "deg2rad", - "diag_indices", - "expand_dims", - "isclose", - "isin", - "kron", - "nan_to_num", - "nanmax", - "nanmin", - "nansum", - "nunique", - "one_hot", - "pad", - "partition", - "rad2deg", - "searchsorted", - "setdiff1d", - "sinc", - "tril_indices", - "triu_indices", - "union1d", - "unravel_index", -] - - -@overload -def apply_where( # numpydoc ignore=GL08 - cond: Array, - args: Array | tuple[Array, ...], - f1: Callable[..., Array], - f2: Callable[..., Array], - /, - *, - kwargs: dict[str, Array] | None = None, - xp: ArrayNamespace | None = None, -) -> Array: ... - - -@overload -def apply_where( # numpydoc ignore=GL08 - cond: Array, - args: Array | tuple[Array, ...], - f1: Callable[..., Array], - /, - *, - fill_value: Array | complex, - kwargs: dict[str, Array] | None = None, - xp: ArrayNamespace | None = None, -) -> Array: ... - - -def apply_where( # numpydoc ignore=PR01,PR02 - cond: Array, - args: Array | tuple[Array, ...], - f1: Callable[..., Array], - f2: Callable[..., Array] | None = None, - /, - *, - fill_value: Array | complex | None = None, - kwargs: dict[str, Array] | None = None, - xp: ArrayNamespace | None = None, -) -> Array: - """ - Run one of two elementwise functions depending on a condition. - - Equivalent to ``f1(*args) if cond else fill_value`` performed elementwise - when `fill_value` is defined, otherwise to ``f1(*args) if cond else f2(*args)``. - - Parameters - ---------- - cond : array - The condition, expressed as a boolean array. - args : Array or tuple of Arrays - Argument(s) to `f1` (and `f2`). Must be broadcastable with `cond`. - f1 : callable - Elementwise function of `args`, returning a single array. - Where `cond` is True, output will be ``f1(arg0[cond], arg1[cond], ...)``. - f2 : callable, optional - Elementwise function of `args`, returning a single array. - Where `cond` is False, output will be ``f2(arg0[cond], arg1[cond], ...)``. - Mutually exclusive with `fill_value`. - fill_value : Array or scalar, optional - If provided, value with which to fill output array where `cond` is False. - It does not need to be scalar; it needs however to be broadcastable with - `cond` and `args`. - Mutually exclusive with `f2`. You must provide one or the other. - kwargs : dict of str : Array pairs - Keyword argument(s) to `f1` (and `f2`). Values must be broadcastable with - `cond`. - xp : array_namespace, optional - The standard-compatible namespace for `cond` and `args`. Default: infer. - - Returns - ------- - Array - An array with elements from the output of `f1` where `cond` is True and either - the output of `f2` or `fill_value` where `cond` is False. The returned array has - data type determined by type promotion rules between the output of `f1` and - either `fill_value` or the output of `f2`. - - Notes - ----- - ``xp.where(cond, f1(*args), f2(*args))`` requires explicitly evaluating `f1` even - when `cond` is False, and `f2` when cond is True. This function evaluates each - function only for their matching condition, if the backend allows for it. - - On Dask, `f1` and `f2` are applied to the individual chunks and should use functions - from the namespace of the chunks. - - Examples - -------- - >>> import array_api_strict as xp - >>> import array_api_extra as xpx - >>> a = xp.asarray([5, 4, 3]) - >>> b = xp.asarray([0, 2, 2]) - >>> def f(a, b): - ... return a // b - >>> xpx.apply_where(b != 0, (a, b), f, fill_value=xp.nan) - array([ nan, 2., 1.]) - """ - # Parse and normalize arguments - if (f2 is None) == (fill_value is None): - msg = "Exactly one of `fill_value` or `f2` must be given." - raise TypeError(msg) - args_ = list(args) if isinstance(args, tuple) else [args] - del args - - kwargs_ = {} if kwargs is None else kwargs - kwkeys = list(kwargs_.keys()) - args_ = [*args_, *kwargs_.values()] - del kwargs - - xp = array_namespace(cond, fill_value, *args_) if xp is None else xp - - if isinstance(fill_value, int | float | complex | NoneType): - cond, *args_ = xp.broadcast_arrays(cond, *args_) - else: - cond, fill_value, *args_ = xp.broadcast_arrays(cond, fill_value, *args_) - - if is_dask_namespace(xp): - meta_xp = meta_namespace(cond, fill_value, *args_, xp=xp) - # map_blocks doesn't descend into tuples of Arrays - return xp.map_blocks( - _apply_where, cond, f1, f2, fill_value, *args_, kwkeys=kwkeys, xp=meta_xp - ) - - return _apply_where(cond, f1, f2, fill_value, *args_, kwkeys=kwkeys, xp=xp) - - -def _apply_where( # numpydoc ignore=PR01,RT01 - cond: Array, - f1: Callable[..., Array], - f2: Callable[..., Array] | None, - fill_value: Array | complex | bool | None, - *args: Array, - kwkeys: list[str], - xp: ArrayNamespace, -) -> Array: - """Helper of `apply_where`. On Dask, this runs on a single chunk.""" - - nargs = len(args) - len(kwkeys) - kwargs = dict(zip(kwkeys, args[nargs:], strict=True)) - args = args[:nargs] - - if not capabilities(xp, device=_compat.device(cond))["boolean indexing"]: - # jax.jit does not support assignment by boolean mask - return xp.where( - cond, - f1(*args, **kwargs), - f2(*args, **kwargs) if f2 is not None else fill_value, - ) - - temp1 = f1( - *(arr[cond] for arr in args), **{key: val[cond] for key, val in kwargs.items()} - ) - - if f2 is None: - dtype = xp.result_type(temp1, fill_value) - if isinstance(fill_value, int | float | complex): - out = xp.full_like(cond, dtype=dtype, fill_value=fill_value) - else: - out = xp.astype(fill_value, dtype, copy=True) - else: - ncond = ~cond - temp2 = f2( - *(arr[ncond] for arr in args), - **{key: val[ncond] for key, val in kwargs.items()}, - ) - dtype = xp.result_type(temp1, temp2) - out = xp.empty_like(cond, dtype=dtype) - out = at(out, ncond).set(temp2) - - return at(out, cond).set(temp1) - - -def atleast_nd(x: Array, /, *, ndim: int, xp: ArrayNamespace) -> Array: - # numpydoc ignore=PR01,RT01 - """See docstring in array_api_extra._delegation.""" - - if x.ndim < ndim: - x = xp.expand_dims(x, axis=0) - x = atleast_nd(x, ndim=ndim, xp=xp) - return x - - -# `float` in signature to accept `math.nan` for Dask. -# `int`s are still accepted as `float` is a superclass of `int` in typing -def broadcast_shapes( # numpydoc ignore=PR01,RT01 - *shapes: tuple[float | None, ...], -) -> tuple[int | None, ...]: - """See docstring in array_api_extra._delegation.""" - if not shapes: - return () # Match NumPy output - - ndim = max(len(shape) for shape in shapes) - out: list[int | None] = [] - for axis in range(-ndim, 0): - sizes = {shape[axis] for shape in shapes if axis >= -len(shape)} - # Dask uses NaN for unknown shape, which predates the Array API spec for None - none_size = None in sizes or math.nan in sizes # noqa: PLW0177 - sizes -= {1, None, math.nan} - if len(sizes) > 1: - msg = ( - "shape mismatch: objects cannot be broadcast to a single shape: " - f"{shapes}." - ) - raise ValueError(msg) - out.append(None if none_size else cast(int, sizes.pop()) if sizes else 1) - - return tuple(out) - - -def cov( - m: Array, - /, - *, - correction: float = 1, - fweights: Array | None = None, - aweights: Array | None = None, - xp: ArrayNamespace, -) -> Array: # numpydoc ignore=PR01,RT01 - """See docstring in array_api_extra._delegation.""" - # NB: no `xp.asarray(m)` here. The delegation layer already guarantees `m` - # is an array (it calls `array_namespace(m)` and reads `m.ndim`), and on - # torch `xp.asarray` detaches gradients and mutates the caller's tensor. - dtype = ( - xp.float64 if xp.isdtype(m.dtype, "integral") else xp.result_type(m, xp.float64) - ) - - m = atleast_nd(m, ndim=2, xp=xp) - # Preserve the historical no-alias guarantee even when the dtype already matches. - m = xp.astype(m, dtype, copy=True) - - # Validate weight shapes (eager metadata, lazy-safe). - n_obs = m.shape[-1] - for name, w_in in (("fweights", fweights), ("aweights", aweights)): - if w_in is None: - continue - if w_in.ndim != 1: - msg = f"`{name}` must be 1-D, got ndim={w_in.ndim}" - raise ValueError(msg) - weight_length = w_in.shape[0] - # Unknown dims are `None` per the standard; Dask non-standardly - # reports them as NaN, hence the `isnan` checks below. - if ( - weight_length is not None - and n_obs is not None - and not math.isnan(weight_length) - and not math.isnan(n_obs) - and weight_length != n_obs - ): - msg = ( - f"`{name}` has length {weight_length} but `m` has {n_obs} observations" - ) - raise ValueError(msg) - - fw = None - if fweights is not None: - fw = xp.astype(xp.asarray(fweights), dtype) - aw = None - if aweights is not None: - aw = xp.astype(xp.asarray(aweights), dtype) - if fw is None and aw is None: - w = None - elif fw is None: - w = aw - elif aw is None: - w = fw - else: - w = fw * aw - - if w is None: - avg = xp.mean(m, axis=-1, keepdims=True) - fact = eager_shape(m, axis=-1)[0] - correction - else: - v1 = xp.sum(w, axis=-1) - avg = xp.sum(m * w, axis=-1, keepdims=True) / v1 - if aw is None: - fact = v1 - correction - else: - fact = v1 - correction * xp.sum(w * aw, axis=-1) / v1 - - if not _compat.is_lazy_array(fact): - # Weights are cast to `dtype`, so a complex input produces a complex - # normalizer with a zero imaginary part. Complex ordering is undefined; - # compare its real component instead. - if w is not None: - fact_array = cast(Array, fact) - fact_to_check = ( - xp.real(fact_array) - if xp.isdtype(fact_array.dtype, "complex floating") - else fact_array - ) - else: - fact_to_check = fact - if fact_to_check <= 0: - warnings.warn( - "Degrees of freedom <= 0 for slice", RuntimeWarning, stacklevel=2 - ) - fact = 0 - - m_c = m - avg - m_w = m_c if w is None else m_c * w - m_cT = xp.matrix_transpose(m_c) - if xp.isdtype(m_cT.dtype, "complex floating"): - m_cT = xp.conj(m_cT) - c = m_w @ m_cT / fact - axes = tuple(axis for axis, length in enumerate(c.shape) if length == 1) - return xp.squeeze(c, axis=axes) - - -def one_hot( - x: Array, - /, - num_classes: int, - *, - xp: ArrayNamespace, -) -> Array: # numpydoc ignore=PR01,RT01 - """See docstring in `array_api_extra._delegation.py`.""" - # TODO: Benchmark whether this is faster on the NumPy backend: - # if is_numpy_array(x): - # out = xp.zeros((x.size, num_classes), dtype=dtype) - # out[xp.arange(x.size), xp.reshape(x, (-1,))] = 1 - # return xp.reshape(out, (*x.shape, num_classes)) - range_num_classes = xp.arange(num_classes, dtype=x.dtype, device=_compat.device(x)) - return x[..., xp.newaxis] == range_num_classes - - -def create_diagonal( - x: Array, /, *, offset: int = 0, xp: ArrayNamespace -) -> Array: # numpydoc ignore=PR01,RT01 - """See docstring in array_api_extra._delegation.""" - x_shape = eager_shape(x) - batch_dims = x_shape[:-1] - n = x_shape[-1] + abs(offset) - diag = xp.zeros((*batch_dims, n**2), dtype=x.dtype, device=_compat.device(x)) - - target_slice = slice( - offset if offset >= 0 else abs(offset) * n, - min(n * (n - offset), diag.shape[-1]), - n + 1, - ) - for index in ndindex(*batch_dims): - diag = at(diag)[(*index, target_slice)].set(x[(*index, slice(None))]) - return xp.reshape(diag, (*batch_dims, n, n)) - - -def diag_indices( - n: int, /, *, ndim: int, device: Device | None, xp: ArrayNamespace -) -> tuple[Array, ...]: # numpydoc ignore=PR01,RT01 - """See docstring in array_api_extra._delegation.""" - idx = xp.arange(n, device=device) - return (idx,) * ndim - - -def _tri_indices( - n: int, - *, - offset: int, - m: int | None, - upper: bool, - device: Device | None, - xp: ArrayNamespace, -) -> tuple[Array, Array]: # numpydoc ignore=PR01,RT01 - """Shared implementation for `tril_indices` and `triu_indices`.""" - cols = n if m is None else m - rows = xp.arange(n, device=device)[:, xp.newaxis] - cols_a = xp.arange(cols, device=device)[xp.newaxis, :] - delta = cols_a - rows - mask = delta >= offset if upper else delta <= offset - r, c = xp.nonzero(mask) - return (r, c) - - -def tril_indices( - n: int, - /, - *, - offset: int, - m: int | None, - device: Device | None, - xp: ArrayNamespace, -) -> tuple[Array, Array]: # numpydoc ignore=PR01,RT01 - """See docstring in array_api_extra._delegation.""" - return _tri_indices(n, offset=offset, m=m, upper=False, device=device, xp=xp) - - -def triu_indices( - n: int, - /, - *, - offset: int, - m: int | None, - device: Device | None, - xp: ArrayNamespace, -) -> tuple[Array, Array]: # numpydoc ignore=PR01,RT01 - """See docstring in array_api_extra._delegation.""" - return _tri_indices(n, offset=offset, m=m, upper=True, device=device, xp=xp) - - -def default_dtype( - xp: ArrayNamespace, - kind: Literal[ - "real floating", "complex floating", "integral", "indexing" - ] = "real floating", - *, - device: Device | None = None, -) -> DType: - """ - Return the default dtype for the given namespace and device. - - This is a convenience shorthand for - ``xp.__array_namespace_info__().default_dtypes(device=device)[kind]``. - - Parameters - ---------- - xp : array_namespace - The standard-compatible namespace for which to get the default dtype. - kind : {'real floating', 'complex floating', 'integral', 'indexing'}, optional - The kind of dtype to return. Default is 'real floating'. - device : Device, optional - The device for which to get the default dtype. Default: current device. - - Returns - ------- - dtype - The default dtype for the given namespace, kind, and device. - """ - dtypes = xp.__array_namespace_info__().default_dtypes(device=device) - try: - return dtypes[kind] - except KeyError as e: - domain = ("real floating", "complex floating", "integral", "indexing") - assert set(dtypes) == set(domain), f"Non-compliant namespace: {dtypes}" - msg = f"Unknown kind '{kind}'. Expected one of {domain}." - raise ValueError(msg) from e - - -def expand_dims( - a: Array, /, *, axis: tuple[int, ...] = (0,), xp: ArrayNamespace -) -> Array: - # numpydoc ignore=PR01,RT01 - """See docstring in array_api_extra._delegation.""" - for i in sorted(axis): - a = xp.expand_dims(a, axis=i) - return a - - -def isclose( - a: Array | complex, - b: Array | complex, - *, - rtol: float = 1e-05, - atol: float = 1e-08, - equal_nan: bool = False, - xp: ArrayNamespace, -) -> Array: # numpydoc ignore=PR01,RT01 - """See docstring in array_api_extra._delegation.""" - a, b = asarrays(a, b, xp=xp) - - a_inexact = xp.isdtype(a.dtype, ("real floating", "complex floating")) - b_inexact = xp.isdtype(b.dtype, ("real floating", "complex floating")) - if a_inexact or b_inexact: - # prevent warnings on NumPy and Dask on inf - inf - mxp = meta_namespace(a, b, xp=xp) - out = apply_where( - xp.isinf(a) | xp.isinf(b), - (a, b), - lambda a, b: mxp.isinf(a) & mxp.isinf(b) & (mxp.sign(a) == mxp.sign(b)), # pyright: ignore[reportUnknownArgumentType] - # Note: inf <= inf is True! - lambda a, b: mxp.abs(a - b) <= (atol + rtol * mxp.abs(b)), # pyright: ignore[reportUnknownArgumentType] - xp=xp, - ) - if equal_nan: - out = xp.where(xp.isnan(a) & xp.isnan(b), True, out) - return out - - if xp.isdtype(a.dtype, "bool") or xp.isdtype(b.dtype, "bool"): - if atol >= 1 or rtol >= 1: - return xp.ones_like(a == b) - return a == b - - # integer types - atol = int(atol) - if rtol == 0: - return xp.abs(a - b) <= atol - - # Don't rely on OverflowError, as it is not guaranteed by the Array API. - nrtol = int(1.0 / rtol) - if nrtol > xp.iinfo(b.dtype).max: - # rtol * max_int < 1, so it's inconsequential - return xp.abs(a - b) <= atol - return xp.abs(a - b) <= (atol + xp.abs(b) // nrtol) - - -def kron( - a: Array, - b: Array, - /, - *, - xp: ArrayNamespace, -) -> Array: # numpydoc ignore=PR01,RT01 - """See docstring in array_api_extra._delegation.""" - - singletons = (1,) * (b.ndim - a.ndim) - a = cast(Array, xp.broadcast_to(a, singletons + a.shape)) - - nd_b, nd_a = b.ndim, a.ndim - nd_max = max(nd_b, nd_a) - if nd_a == 0 or nd_b == 0: - return xp.multiply(a, b) - - a_shape = eager_shape(a) - b_shape = eager_shape(b) - - # Equalise the shapes by prepending smaller one with 1s - a_shape = (1,) * max(0, nd_b - nd_a) + a_shape - b_shape = (1,) * max(0, nd_a - nd_b) + b_shape - - # Insert empty dimensions - a_arr = expand_dims(a, axis=tuple(range(nd_b - nd_a)), xp=xp) - b_arr = expand_dims(b, axis=tuple(range(nd_a - nd_b)), xp=xp) - - # Compute the product - a_arr = expand_dims(a_arr, axis=tuple(range(1, nd_max * 2, 2)), xp=xp) - b_arr = expand_dims(b_arr, axis=tuple(range(0, nd_max * 2, 2)), xp=xp) - result = xp.multiply(a_arr, b_arr) - - # Reshape back and return - res_shape = tuple(a_s * b_s for a_s, b_s in zip(a_shape, b_shape, strict=True)) - return xp.reshape(result, res_shape) - - -def nan_to_num( # numpydoc ignore=PR01,RT01 - x: Array, - /, - fill_value: float = 0.0, - *, - xp: ArrayNamespace, -) -> Array: - """See docstring in `array_api_extra._delegation.py`.""" - - def perform_replacements( # numpydoc ignore=PR01,RT01 - x: Array, - fill_value: float, - xp: ArrayNamespace, - ) -> Array: - """Internal function to perform the replacements.""" - x = xp.where(xp.isnan(x), fill_value, x) - - # convert infinities to finite values - finfo = xp.finfo(x.dtype) - idx_posinf = xp.isinf(x) & ~xp.signbit(x) - idx_neginf = xp.isinf(x) & xp.signbit(x) - x = xp.where(idx_posinf, finfo.max, x) - return xp.where(idx_neginf, finfo.min, x) - - if xp.isdtype(x.dtype, "complex floating"): - return perform_replacements( - xp.real(x), - fill_value, - xp, - ) + 1j * perform_replacements( - xp.imag(x), - fill_value, - xp, - ) - - if xp.isdtype(x.dtype, "numeric"): - return perform_replacements(x, fill_value, xp) - - return x - - -def nunique(x: Array, /, *, xp: ArrayNamespace) -> Array: # numpydoc ignore=PR01,RT01 - """See docstring in `array_api_extra._delegation.py`.""" - # There are 3 general use cases: - # 1. backend has unique_counts and it returns an array with known shape - # 2. backend has unique_counts and it returns a None-sized array; - # e.g. Dask, ndonnx - # 3. backend does not have unique_counts; e.g. wrapped JAX - if capabilities(xp, device=_compat.device(x))["data-dependent shapes"]: - # xp has unique_counts; O(n) complexity - _, counts = xp.unique_counts(x) - n = _compat.size(counts) - if n is None: - return xp.sum(xp.ones_like(counts)) - return xp.asarray(n, device=_compat.device(x)) - - # xp does not have unique_counts; O(n*logn) complexity - x = xp.reshape(x, (-1,)) - x = xp.sort(x, stable=False) - mask = x != xp.roll(x, -1) - default_int = default_dtype(xp, "integral", device=_compat.device(x)) - return xp.maximum( - # Special cases: - # - array is size 0 - # - array has all elements equal to each other - xp.astype(xp.any(~mask), default_int), - xp.sum(xp.astype(mask, default_int)), - ) - - -def pad( - x: Array, - pad_width: int | tuple[int, int] | Sequence[tuple[int, int]], - *, - constant_values: complex = 0, - xp: ArrayNamespace, -) -> Array: # numpydoc ignore=PR01,RT01 - """See docstring in `array_api_extra._delegation.py`.""" - pad_width_seq = normalize_pad_width(pad_width, x.ndim) - - slices: list[slice] = [] - newshape: list[int] = [] - for ax, w_tpl in enumerate(pad_width_seq): - if len(w_tpl) != 2: - msg = f"expect a 2-tuple (before, after), got {w_tpl}." - raise ValueError(msg) - - sh = eager_shape(x)[ax] - - if w_tpl[0] == 0 and w_tpl[1] == 0: - sl = slice(None, None, None) - else: - stop: int | None - start, stop = w_tpl - stop = None if stop == 0 else -stop - - sl = slice(start, stop, None) - sh += w_tpl[0] + w_tpl[1] - - newshape.append(sh) - slices.append(sl) - - padded = xp.full( - tuple(newshape), - fill_value=constant_values, - dtype=x.dtype, - device=_compat.device(x), - ) - return at(padded, tuple(slices)).set(x) - - -def searchsorted( - x1: Array, - x2: Array, - /, - *, - side: Literal["left", "right"] = "left", - xp: ArrayNamespace, -) -> Array: - # numpydoc ignore=PR01,RT01 - """See docstring in `array_api_extra._delegation.py`.""" - a = xp.full(x2.shape, 0, device=_compat.device(x1)) - - if x1.shape[-1] == 0: - return a - - n = xp.count_nonzero(~xp.isnan(x1), axis=-1, keepdims=True) - b = xp.broadcast_to(n, x2.shape) - - compare = xp.less_equal if side == "left" else xp.less - - # while xp.any(b - a > 1): - # refactored to for loop with ~log2(n) iterations for JAX JIT - for _ in range(int(math.log2(x1.shape[-1])) + 1): # type: ignore[arg-type] # pyright: ignore[reportArgumentType] - c = (a + b) // 2 - x0 = xp.take_along_axis(x1, c, axis=-1) - j = compare(x2, x0) - b = xp.where(j, c, b) - a = xp.where(j, a, c) - - out = xp.where(compare(x2, xp.min(x1, axis=-1, keepdims=True)), 0, b) - out = xp.where(xp.isnan(x2), x1.shape[-1], out) if side == "right" else out - return xp.astype(out, default_dtype(xp, kind="integral"), copy=False) - - -def setdiff1d( - x1: Array | complex, - x2: Array | complex, - /, - *, - assume_unique: bool = False, - xp: ArrayNamespace, -) -> Array: # numpydoc ignore=PR01,RT01 - """See docstring in `array_api_extra._delegation.py`.""" - - # https://github.com/microsoft/pyright/issues/10103 - x1_, x2_ = asarrays(x1, x2, xp=xp) - - if assume_unique: - x1_ = xp.reshape(x1_, (-1,)) - x2_ = xp.reshape(x2_, (-1,)) - else: - x1_ = xp.unique_values(x1_) - x2_ = xp.unique_values(x2_) - - return x1_[_helpers.in1d(x1_, x2_, assume_unique=True, invert=True, xp=xp)] - - -def sinc(x: Array, /, *, xp: ArrayNamespace) -> Array: - # numpydoc ignore=PR01,RT01 - """See docstring in `array_api_extra._delegation.py`.""" - - # no scalars in `where` - array-api#807 - y = xp.pi * xp.where( - xp.astype(x, xp.bool), - x, - xp.asarray(xp.finfo(x.dtype).eps, dtype=x.dtype, device=_compat.device(x)), - ) - return xp.sin(y) / y - - -def partition( # numpydoc ignore=PR01,RT01 - x: Array, - kth: int, # noqa: ARG001 - /, - axis: int = -1, - *, - xp: ArrayNamespace, -) -> Array: - """See docstring in `array_api_extra._delegation.py`.""" - return xp.sort(x, axis=axis, stable=False) - - -def argpartition( # numpydoc ignore=PR01,RT01 - x: Array, - kth: int, # noqa: ARG001 - /, - axis: int = -1, - *, - xp: ArrayNamespace, -) -> Array: - """See docstring in `array_api_extra._delegation.py`.""" - return xp.argsort(x, axis=axis, stable=False) - - -def isin( # numpydoc ignore=PR01,RT01 - a: Array, - b: Array, - /, - *, - assume_unique: bool = False, - invert: bool = False, - xp: ArrayNamespace, -) -> Array: - """See docstring in `array_api_extra._delegation.py`.""" - original_a_shape = a.shape - a = xp.reshape(a, (-1,)) - b = xp.reshape(b, (-1,)) - return xp.reshape( - _helpers.in1d(a, b, assume_unique=assume_unique, invert=invert, xp=xp), - original_a_shape, - ) - - -def union1d(a: Array, b: Array, /, *, xp: ArrayNamespace) -> Array: - # numpydoc ignore=PR01,RT01 - """See docstring in `array_api_extra._delegation.py`.""" - a = xp.reshape(a, (-1,)) - b = xp.reshape(b, (-1,)) - # XXX: `sparse` returns NumPy arrays from `unique_values` - return xp.asarray(xp.unique_values(xp.concat([a, b]))) - - -def angle(z: Array, /, *, deg: bool = False, xp: ArrayNamespace | None = None) -> Array: - """ - Return the angle of the complex argument. - - Parameters - ---------- - z : Array - Input array. - deg : bool, optional - Return angle in degrees if True, radians if False (default). - xp : array_namespace, optional - The standard-compatible namespace for `z`. Default: infer. - - Returns - ------- - array - The counterclockwise angle from the positive real axis on the complex - plane in the range ``(-pi, pi]``. - - Notes - ----- - Real input ``x`` is interpreted as ``x + 0j``. - - Examples - -------- - >>> import array_api_strict as xp - >>> import array_api_extra as xpx - >>> xpx.angle(xp.asarray([1.0, 1.0j, 1 + 1j]), xp=xp) - Array([0. , 1.57079633, 0.78539816], dtype=array_api_strict.float64) - >>> xpx.angle(xp.asarray([1.0, 1.0j, 1 + 1j]), deg=True, xp=xp) - Array([ 0., 90., 45.], dtype=array_api_strict.float64) - """ - if xp is None: - xp = array_namespace(z) - if xp.isdtype(z.dtype, "complex floating"): - zimag = xp.imag(z) - zreal = xp.real(z) - else: - if not xp.isdtype(z.dtype, "real floating"): - z = xp.astype(z, default_dtype(xp, device=_compat.device(z))) - zimag = xp.zeros_like(z) - zreal = z - a = xp.atan2(zimag, zreal) - if deg: - a = a * 180 / xp.pi - return a - - -def deg2rad(x: Array, /, *, xp: ArrayNamespace) -> Array: - # numpydoc ignore=PR01,RT01 - """See docstring in `array_api_extra._delegation.py`.""" - return x * xp.pi / 180 - - -def rad2deg(x: Array, /, *, xp: ArrayNamespace) -> Array: - # numpydoc ignore=PR01,RT01 - """See docstring in `array_api_extra._delegation.py`.""" - return x * 180 / xp.pi - - -def unravel_index(indices: Array, shape: tuple[int, ...], /) -> tuple[Array, ...]: - # numpydoc ignore=PR01,RT01 - """See docstring in `array_api_extra._delegation.py`.""" - coords: list[Array] = [] - for dim in reversed(shape): - coords.append(indices % dim) - indices = indices // dim - return tuple(reversed(coords)) - - -def nanmin( # numpydoc ignore=PR01,RT01 - a: Array, - /, - *, - axis: int | tuple[int, ...] | None, - xp: ArrayNamespace, -) -> Array: - """See docstring in `array_api_extra._delegation.py`.""" - mask = xp.isnan(a) - device_a = _compat.device(a) - x = xp.min( - xp.where(mask, xp.asarray(+xp.inf, dtype=a.dtype, device=device_a), a), - axis=axis, - ) - # Replace Infs from all NaN slices with NaN again - mask = xp.all(mask, axis=axis) - if xp.any(mask): - x = xp.where(mask, xp.asarray(xp.nan, dtype=x.dtype, device=device_a), x) - return x - - -def nanmax( # numpydoc ignore=PR01,RT01 - a: Array, - /, - *, - axis: int | tuple[int, ...] | None, - xp: ArrayNamespace, -) -> Array: - """See docstring in `array_api_extra._delegation.py`.""" - mask = xp.isnan(a) - device_a = _compat.device(a) - x = xp.max( - xp.where(mask, xp.asarray(-xp.inf, dtype=a.dtype, device=device_a), a), - axis=axis, - ) - # Replace Infs from all NaN slices with NaN again - mask = xp.all(mask, axis=axis) - if xp.any(mask): - x = xp.where(mask, xp.asarray(xp.nan, dtype=x.dtype, device=device_a), x) - return x - - -def nansum( # numpydoc ignore=PR01,RT01 - a: Array, - /, - *, - axis: int | tuple[int, ...] | None, - xp: ArrayNamespace, -) -> Array: - """See docstring in `array_api_extra._delegation.py`.""" - mask = xp.isnan(a) - device_a = _compat.device(a) - zero = xp.asarray(0, dtype=a.dtype, device=device_a) - return xp.sum(xp.where(mask, zero, a), axis=axis) diff --git a/src/array_api_extra/_lib/_utils/_helpers.py b/src/array_api_extra/_lib/_helpers.py similarity index 99% rename from src/array_api_extra/_lib/_utils/_helpers.py rename to src/array_api_extra/_lib/_helpers.py index 5cb3c267..5d9712e6 100644 --- a/src/array_api_extra/_lib/_utils/_helpers.py +++ b/src/array_api_extra/_lib/_helpers.py @@ -1,4 +1,4 @@ -"""Helper functions used by `array_api_extra/_funcs.py`.""" +"""Helper functions.""" from __future__ import annotations diff --git a/src/array_api_extra/_lib/_testing.py b/src/array_api_extra/_lib/_testing.py index 6dfa0fc2..9dbc3943 100644 --- a/src/array_api_extra/_lib/_testing.py +++ b/src/array_api_extra/_lib/_testing.py @@ -1,8 +1,4 @@ -""" -Private testing utilities. - -See also ..testing for public testing utilities. -""" +"""Private testing utilities.""" from __future__ import annotations diff --git a/src/array_api_extra/_lib/_utils/_typing.py b/src/array_api_extra/_lib/_typing.py similarity index 79% rename from src/array_api_extra/_lib/_utils/_typing.py rename to src/array_api_extra/_lib/_typing.py index 651b38df..388c3534 100644 --- a/src/array_api_extra/_lib/_utils/_typing.py +++ b/src/array_api_extra/_lib/_typing.py @@ -1,5 +1,6 @@ +"""Static typing helpers.""" # numpydoc ignore=GL08 -# pylint: disable=missing-module-docstring,duplicate-code +# pylint: disable=duplicate-code from types import ModuleType diff --git a/src/array_api_extra/_lib/_utils/_typing.pyi b/src/array_api_extra/_lib/_typing.pyi similarity index 100% rename from src/array_api_extra/_lib/_utils/_typing.pyi rename to src/array_api_extra/_lib/_typing.pyi diff --git a/src/array_api_extra/_lib/_utils/__init__.py b/src/array_api_extra/_lib/_utils/__init__.py deleted file mode 100644 index 3628c45f..00000000 --- a/src/array_api_extra/_lib/_utils/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Modules housing private utility functions.""" diff --git a/src/array_api_extra/_linalg.py b/src/array_api_extra/_linalg.py new file mode 100644 index 00000000..5ec436b4 --- /dev/null +++ b/src/array_api_extra/_linalg.py @@ -0,0 +1,109 @@ +"""Delegation layer for linear algebra functions.""" + +from ._agnostic import _linalg +from ._lib._compat import ( + array_namespace, + is_cupy_namespace, + is_jax_namespace, + is_numpy_namespace, + is_torch_namespace, +) +from ._lib._helpers import asarrays +from ._lib._typing import Array, ArrayNamespace + +__all__ = ["kron"] + + +def kron( + a: Array | complex, + b: Array | complex, + /, + *, + xp: ArrayNamespace | None = None, +) -> Array: + """ + Kronecker product of two arrays. + + Computes the Kronecker product, a composite array made of blocks of the + second array scaled by the first. + + Equivalent to ``numpy.kron`` for NumPy arrays. + + Parameters + ---------- + a, b : Array | int | float | complex + Input arrays or scalars. At least one must be an array. + xp : array_namespace, optional + The standard-compatible namespace for `a` and `b`. Default: infer. + + Returns + ------- + array + The Kronecker product of `a` and `b`. + + Notes + ----- + The function assumes that the number of dimensions of `a` and `b` + are the same, if necessary prepending the smallest with ones. + If ``a.shape = (r0,r1,..,rN)`` and ``b.shape = (s0,s1,...,sN)``, + the Kronecker product has shape ``(r0*s0, r1*s1, ..., rN*SN)``. + The elements are products of elements from `a` and `b`, organized + explicitly by:: + + kron(a,b)[k0,k1,...,kN] = a[i0,i1,...,iN] * b[j0,j1,...,jN] + + where:: + + kt = it * st + jt, t = 0,...,N + + In the common 2-D case (N=1), the block structure can be visualized:: + + [[ a[0,0]*b, a[0,1]*b, ... , a[0,-1]*b ], + [ ... ... ], + [ a[-1,0]*b, a[-1,1]*b, ... , a[-1,-1]*b ]] + + Examples + -------- + >>> import array_api_strict as xp + >>> import array_api_extra as xpx + >>> xpx.kron(xp.asarray([1, 10, 100]), xp.asarray([5, 6, 7]), xp=xp) + Array([ 5, 6, 7, 50, 60, 70, 500, + 600, 700], dtype=array_api_strict.int64) + + >>> xpx.kron(xp.asarray([5, 6, 7]), xp.asarray([1, 10, 100]), xp=xp) + Array([ 5, 50, 500, 6, 60, 600, 7, + 70, 700], dtype=array_api_strict.int64) + + >>> xpx.kron(xp.eye(2), xp.ones((2, 2)), xp=xp) + Array([[1., 1., 0., 0.], + [1., 1., 0., 0.], + [0., 0., 1., 1.], + [0., 0., 1., 1.]], dtype=array_api_strict.float64) + + >>> a = xp.reshape(xp.arange(100), (2, 5, 2, 5)) + >>> b = xp.reshape(xp.arange(24), (2, 3, 4)) + >>> c = xpx.kron(a, b, xp=xp) + >>> c.shape + (2, 10, 6, 20) + >>> I = (1, 3, 0, 2) + >>> J = (0, 2, 1) + >>> J1 = (0,) + J # extend to ndim=4 + >>> S1 = (1,) + b.shape + >>> K = tuple(xp.asarray(I) * xp.asarray(S1) + xp.asarray(J1)) + >>> c[K] == a[I]*b[J] + Array(True, dtype=array_api_strict.bool) + """ + if xp is None: + xp = array_namespace(a, b) + + a, b = asarrays(a, b, xp=xp) + + if ( + is_cupy_namespace(xp) + or is_jax_namespace(xp) + or is_numpy_namespace(xp) + or is_torch_namespace(xp) + ): + return xp.kron(a, b) + + return _linalg.kron(a, b, xp=xp) diff --git a/src/array_api_extra/_manipulation.py b/src/array_api_extra/_manipulation.py new file mode 100644 index 00000000..132d41ec --- /dev/null +++ b/src/array_api_extra/_manipulation.py @@ -0,0 +1,287 @@ +"""Delegation layer for manipulation functions.""" + +from collections.abc import Sequence +from typing import Literal + +from ._agnostic import _manipulation +from ._lib._compat import ( + array_namespace, + is_cupy_namespace, + is_dask_namespace, + is_jax_namespace, + is_numpy_namespace, + is_pydata_sparse_namespace, + is_torch_namespace, +) +from ._lib._helpers import deprecated, normalize_pad_width +from ._lib._typing import Array, ArrayNamespace + +__all__ = ["atleast_nd", "broadcast_shapes", "expand_dims", "pad"] + + +def atleast_nd(x: Array, /, *, ndim: int, xp: ArrayNamespace | None = None) -> Array: + """ + Recursively expand the dimension of an array to at least `ndim`. + + Parameters + ---------- + x : array + Input array. + ndim : int + The minimum number of dimensions for the result. + xp : array_namespace, optional + The standard-compatible namespace for `x`. Default: infer. + + Returns + ------- + array + An array with ``res.ndim`` >= `ndim`. + If ``x.ndim`` >= `ndim`, `x` is returned. + If ``x.ndim`` < `ndim`, `x` is expanded by prepending new axes + until ``res.ndim`` equals `ndim`. + + Examples + -------- + >>> import array_api_strict as xp + >>> import array_api_extra as xpx + >>> x = xp.asarray([1]) + >>> xpx.atleast_nd(x, ndim=3, xp=xp) + Array([[[1]]], dtype=array_api_strict.int64) + + >>> x = xp.asarray([[[1, 2], + ... [3, 4]]]) + >>> xpx.atleast_nd(x, ndim=1, xp=xp) is x + True + """ + if xp is None: + xp = array_namespace(x) + + if 1 <= ndim <= 2 and ( + is_numpy_namespace(xp) + or is_jax_namespace(xp) + or is_dask_namespace(xp) + or is_cupy_namespace(xp) + or is_torch_namespace(xp) + ): + return getattr(xp, f"atleast_{ndim}d")(x) + + return _manipulation.atleast_nd(x, ndim=ndim, xp=xp) + + +@deprecated( + "`xpx.broadcast_shapes` is deprecated and will be removed in v1.0.0. " + "`xp.broadcast_shapes` exists in the standard as of v2025.12." +) +def broadcast_shapes( + *shapes: tuple[float | None, ...], xp: ArrayNamespace | None = None +) -> tuple[int | None, ...]: + """ + Compute the shape of the broadcasted arrays. + + .. deprecated:: 0.11.0 + :func:`broadcast_shapes` is deprecated and will be removed in v1.0.0. + :func:`array_api.broadcast_shapes` exists in the standard as of v2025.12. + + Duplicates :func:`numpy.broadcast_shapes`, with additional support for + None and NaN sizes. + + Parameters + ---------- + *shapes : tuple[int | None, ...] + Shapes of the arrays to broadcast. + xp : array_namespace, optional + The standard-compatible namespace to use for native delegation. + Default: use the array-agnostic implementation. + + Returns + ------- + tuple[int | None, ...] + The shape of the broadcasted arrays. + + See Also + -------- + numpy.broadcast_shapes : Equivalent NumPy function. + array_api.broadcast_arrays : Function to broadcast actual arrays. + + Notes + ----- + This function accepts the Array API's ``None`` for unknown sizes, + as well as Dask's non-standard ``math.nan``. + Regardless of input, the output always contains ``None`` for unknown sizes. + + Examples + -------- + >>> import array_api_extra as xpx + >>> xpx.broadcast_shapes((2, 3), (2, 1)) + (2, 3) + >>> xpx.broadcast_shapes((4, 2, 3), (2, 1), (1, 3)) + (4, 2, 3) + """ + if ( + xp is not None + and all(isinstance(size, int) for shape in shapes for size in shape) + and ( + is_numpy_namespace(xp) + or is_cupy_namespace(xp) + or is_jax_namespace(xp) + or is_torch_namespace(xp) + ) + ): + return xp.broadcast_shapes(*shapes) + + return _manipulation.broadcast_shapes(*shapes) + + +@deprecated( + "`xpx.expand_dims` is deprecated and will be removed in v1.0.0. " + "`xp.expand_dims` with support for a tuple of ints in `axis` " + "exists in the standard as of v2025.12." +) +def expand_dims( + a: Array, /, *, axis: int | tuple[int, ...] = (0,), xp: ArrayNamespace | None = None +) -> Array: + """ + Expand the shape of an array. + + .. deprecated:: 0.11.0 + :func:`expand_dims` is deprecated and will be removed in v1.0.0. + :func:`array_api.expand_dims` with support for a tuple of ints in `axis` + exists in the standard as of v2025.12. + + Insert (a) new axis/axes that will appear at the position(s) specified by + `axis` in the expanded array shape. + + Parameters + ---------- + a : array + Array to have its shape expanded. + axis : int or tuple of ints, optional + Position(s) in the expanded axes where the new axis (or axes) is/are placed. + If multiple positions are provided, they should be unique (note that a position + given by a positive index could also be referred to by a negative index - + that will also result in an error). + Default: ``(0,)``. + xp : array_namespace, optional + The standard-compatible namespace for `a`. Default: infer. + + Returns + ------- + array + `a` with an expanded shape. + + Examples + -------- + >>> import array_api_strict as xp + >>> import array_api_extra as xpx + >>> x = xp.asarray([1, 2]) + >>> x.shape + (2,) + + The following is equivalent to ``x[xp.newaxis, :]`` or ``x[xp.newaxis]``: + + >>> y = xpx.expand_dims(x, axis=0, xp=xp) + >>> y + Array([[1, 2]], dtype=array_api_strict.int64) + >>> y.shape + (1, 2) + + The following is equivalent to ``x[:, xp.newaxis]``: + + >>> y = xpx.expand_dims(x, axis=1, xp=xp) + >>> y + Array([[1], + [2]], dtype=array_api_strict.int64) + >>> y.shape + (2, 1) + + ``axis`` may also be a tuple: + + >>> y = xpx.expand_dims(x, axis=(0, 1), xp=xp) + >>> y + Array([[[1, 2]]], dtype=array_api_strict.int64) + + >>> y = xpx.expand_dims(x, axis=(2, 0), xp=xp) + >>> y + Array([[[1], + [2]]], dtype=array_api_strict.int64) + """ + if xp is None: + xp = array_namespace(a) + + if not isinstance(axis, tuple): + axis = (axis,) + ndim = a.ndim + len(axis) + if axis != () and (min(axis) < -ndim or max(axis) >= ndim): + err_msg = ( + f"a provided axis position is out of bounds for array of dimension {a.ndim}" + ) + raise IndexError(err_msg) + axis = tuple(dim % ndim for dim in axis) + if len(set(axis)) != len(axis): + err_msg = "Duplicate dimensions specified in `axis`." + raise ValueError(err_msg) + + if is_numpy_namespace(xp) or is_dask_namespace(xp) or is_jax_namespace(xp): + return xp.expand_dims(a, axis=axis) + + return _manipulation.expand_dims(a, axis=axis, xp=xp) + + +def pad( + x: Array, + pad_width: int | tuple[int, int] | Sequence[tuple[int, int]], + mode: Literal["constant"] = "constant", + *, + constant_values: complex = 0, + xp: ArrayNamespace | None = None, +) -> Array: + """ + Pad the input array. + + Parameters + ---------- + x : array + Input array. + pad_width : int or tuple of ints or sequence of pairs of ints + Pad the input array with this many elements from each side. + If a sequence of tuples, ``[(before_0, after_0), ... (before_N, after_N)]``, + each pair applies to the corresponding axis of ``x``. + A single tuple, ``(before, after)``, is equivalent to a list of ``x.ndim`` + copies of this tuple. + mode : str, optional + Only "constant" mode is currently supported, which pads with + the value passed to `constant_values`. + constant_values : python scalar, optional + Use this value to pad the input. Default is zero. + xp : array_namespace, optional + The standard-compatible namespace for `x`. Default: infer. + + Returns + ------- + array + The input array, + padded with ``pad_width`` elements equal to ``constant_values``. + """ + xp = array_namespace(x) if xp is None else xp + + if mode != "constant": + msg = "Only `'constant'` mode is currently supported" + raise NotImplementedError(msg) + + if ( + is_numpy_namespace(xp) + or is_cupy_namespace(xp) + or is_jax_namespace(xp) + or is_pydata_sparse_namespace(xp) + ): + return xp.pad(x, pad_width, mode, constant_values=constant_values) + + if is_torch_namespace(xp): + # normalize `pad_width` on the host rather than through a tensor as done in + # `torch/_numpy`'s implementation (avoids device transfers) + pad_width_seq = normalize_pad_width(pad_width, x.ndim) + # torch.nn.functional.pad counts dimensions from the last one + flat_pad_width = [w for pair in reversed(pad_width_seq) for w in pair] + return xp.nn.functional.pad(x, tuple(flat_pad_width), value=constant_values) + + return _manipulation.pad(x, pad_width, constant_values=constant_values, xp=xp) diff --git a/src/array_api_extra/_searching.py b/src/array_api_extra/_searching.py new file mode 100644 index 00000000..0689e078 --- /dev/null +++ b/src/array_api_extra/_searching.py @@ -0,0 +1,88 @@ +"""Delegation layer for searching functions.""" + +from typing import Literal + +from ._agnostic import _inspection, _searching +from ._lib._compat import array_namespace, is_torch_namespace +from ._lib._typing import Array, ArrayNamespace + +__all__ = ["searchsorted"] + + +def searchsorted( + x1: Array, + x2: Array, + /, + *, + side: Literal["left", "right"] = "left", + xp: ArrayNamespace | None = None, +) -> Array: + """ + Find indices where elements should be inserted to maintain order. + + Find the indices into a sorted array ``x1`` such that if the elements in ``x2`` + were inserted before the indices, the resulting array would remain sorted. + + The behavior of this function is similar to that of :func:`array_api.searchsorted`, + but it relaxes the requirement that `x1` must be one-dimensional. + This function is vectorized, treating slices along the last axis + as elements and preceding axes as batch (or "loop") dimensions. + + Parameters + ---------- + x1 : Array + Input array. Should have a real-valued data type. Must be sorted in ascending + order along the last axis. + x2 : Array + Array containing search values. Should have a real-valued data type. Must have + the same shape as ``x1`` except along the last axis. + side : {'left', 'right'}, optional + Argument controlling which index is returned if an element of ``x2`` is equal to + one or more elements of ``x1``: ``'left'`` returns the index of the first of + these elements; ``'right'`` returns the next index after the last of these + elements. Default: ``'left'``. + xp : array_namespace, optional + The standard-compatible namespace for the array arguments. Default: infer. + + Returns + ------- + Array: integer array + An array of indices with the same shape as ``x2``. + + Examples + -------- + >>> import array_api_strict as xp + >>> import array_api_extra as xpx + >>> x = xp.asarray([11, 12, 13, 13, 14, 15]) + >>> xpx.searchsorted(x, xp.asarray([10, 11.5, 14.5, 16]), xp=xp) + Array([0, 1, 5, 6], dtype=array_api_strict.int64) + >>> xpx.searchsorted(x, xp.asarray(13), xp=xp) + Array(2, dtype=array_api_strict.int64) + >>> xpx.searchsorted(x, xp.asarray(13), side='right', xp=xp) + Array(4, dtype=array_api_strict.int64) + + `searchsorted` is vectorized along the last axis. + + >>> x1 = xp.asarray([[1., 2., 3., 4.], [5., 6., 7., 8.]]) + >>> x2 = xp.asarray([[1.1, 3.3], [6.6, 8.8]]) + >>> xpx.searchsorted(x1, x2, xp=xp) + Array([[1, 3], + [2, 4]], dtype=array_api_strict.int64) + """ + if xp is None: + xp = array_namespace(x1, x2) + + if side not in {"left", "right"}: + message = "`side` must be either 'left' or 'right'." + raise ValueError(message) + + xp_default_int = _inspection.default_dtype(xp, kind="integral") + x2_0d = x2.ndim == 0 + x1_1d = x1.ndim <= 1 + + if x1_1d or is_torch_namespace(xp): + x2 = xp.reshape(x2, ()) if (x2_0d and x1_1d) else x2 + out = xp.searchsorted(x1, x2, side=side) + return xp.astype(out, xp_default_int, copy=False) + + return _searching.searchsorted(x1, x2, side=side, xp=xp) diff --git a/src/array_api_extra/_set.py b/src/array_api_extra/_set.py new file mode 100644 index 00000000..99d67432 --- /dev/null +++ b/src/array_api_extra/_set.py @@ -0,0 +1,215 @@ +"""Delegation layer for set functions.""" + +from ._agnostic import _set +from ._lib._compat import ( + array_namespace, + is_cupy_namespace, + is_dask_namespace, + is_jax_array, + is_jax_namespace, + is_numpy_namespace, + is_torch_namespace, + size, +) +from ._lib._compat import device as get_device +from ._lib._helpers import asarrays, capabilities +from ._lib._typing import Array, ArrayNamespace + +__all__ = ["isin", "nunique", "setdiff1d", "union1d"] + + +def isin( + a: Array, + b: Array, + /, + *, + assume_unique: bool = False, + invert: bool = False, + kind: str | None = None, + xp: ArrayNamespace | None = None, +) -> Array: + """ + Determine whether each element in `a` is present in `b`. + + This is :func:`array_api.isin`, with additional `assume_unique` + and `kind` parameters. + + Parameters + ---------- + a : array + Input elements. + b : array + The elements against which to test each element of `a`. + assume_unique : bool, optional + If True, the input arrays are both assumed to be unique which can speed + up the calculation. Default: False. + invert : bool, optional + If True, the values in the returned array are inverted. Default: False. + kind : str | None, optional + The algorithm or method to use. This will not affect the final result, + but will affect the speed and memory use. + For NumPy the options are {None, "sort", "table"}. + For Jax the mapped parameter is instead `method` and the options are + {"compare_all", "binary_search", "sort", and "auto" (default)} + For CuPy, Dask, Torch and the default case this parameter is not present and + thus ignored. Default: None. + xp : array_namespace, optional + The standard-compatible namespace for `a` and `b`. Default: infer. + + Returns + ------- + array + An array having the same shape as that of `a` that is True for elements + that are in `b` and False otherwise. + """ + if xp is None: + xp = array_namespace(a, b) + + if is_numpy_namespace(xp): + return xp.isin(a, b, assume_unique=assume_unique, invert=invert, kind=kind) + if is_jax_namespace(xp): + if kind is None: + kind = "auto" + return xp.isin(a, b, assume_unique=assume_unique, invert=invert, method=kind) + if is_cupy_namespace(xp) or is_torch_namespace(xp) or is_dask_namespace(xp): + return xp.isin(a, b, assume_unique=assume_unique, invert=invert) + + return _set.isin(a, b, assume_unique=assume_unique, invert=invert, xp=xp) + + +def nunique(x: Array, /, *, xp: ArrayNamespace | None = None) -> Array: + """ + Count the number of unique elements in an array. + + Compatible with JAX and Dask, whose laziness would be otherwise + problematic. + + Parameters + ---------- + x : Array + Input array. + xp : array_namespace, optional + The standard-compatible namespace for `x`. Default: infer. + + Returns + ------- + array: 0-dimensional integer array + The number of unique elements in `x`. It can be lazy. + """ + if xp is None: + xp = array_namespace(x) + + if is_jax_array(x): + # size= is JAX-specific + # https://github.com/data-apis/array-api/issues/883 + _, counts = xp.unique_counts(x, size=size(x)) + return (counts > 0).sum() + + if ( + is_numpy_namespace(xp) + or is_cupy_namespace(xp) + or ( + is_torch_namespace(xp) + and capabilities(xp, device=get_device(x))["data-dependent shapes"] + ) + ): + _, counts = xp.unique_counts(x) + return xp.asarray(size(counts), device=get_device(x)) + + return _set.nunique(x, xp=xp) + + +def setdiff1d( + x1: Array | complex, + x2: Array | complex, + /, + *, + assume_unique: bool = False, + xp: ArrayNamespace | None = None, +) -> Array: + """ + Find the set difference of two arrays. + + Return the unique values in `x1` that are not in `x2`. + + Parameters + ---------- + x1 : array | int | float | complex | bool + 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, optional + The standard-compatible namespace for `x1` and `x2`. Default: infer. + + Returns + ------- + 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 xp is None: + xp = array_namespace(x1, x2) + + if is_numpy_namespace(xp) or is_cupy_namespace(xp) or is_jax_namespace(xp): + x1, x2 = asarrays(x1, x2, xp=xp) + return xp.setdiff1d(x1, x2, assume_unique=assume_unique) + + return _set.setdiff1d(x1, x2, assume_unique=assume_unique, xp=xp) + + +def union1d(a: Array, b: Array, /, *, xp: ArrayNamespace | None = None) -> Array: + """ + Find the union of two arrays. + + Return the unique, sorted array of values that are in either of the two + input arrays. + + Parameters + ---------- + a, b : Array + Input arrays. They are flattened internally if they are not already 1D. + + xp : array_namespace, optional + The standard-compatible namespace for `a` and `b`. Default: infer. + + Returns + ------- + Array + Unique, sorted union of the input arrays. + + See Also + -------- + jax.numpy.union1d : Corresponding function in JAX. + + Notes + ----- + This function is not compatible with `jax.jit`. + See the docstring of the corresponding JAX function for more information. + """ + if xp is None: + xp = array_namespace(a, b) + + if ( + is_numpy_namespace(xp) + or is_cupy_namespace(xp) + or is_dask_namespace(xp) + or is_jax_namespace(xp) + ): + return xp.union1d(a, b) + + return _set.union1d(a, b, xp=xp) diff --git a/src/array_api_extra/_sorting.py b/src/array_api_extra/_sorting.py new file mode 100644 index 00000000..b98ed320 --- /dev/null +++ b/src/array_api_extra/_sorting.py @@ -0,0 +1,206 @@ +"""Delegation layer for sorting functions.""" + +from ._agnostic import _sorting +from ._lib._compat import ( + array_namespace, + is_cupy_namespace, + is_jax_namespace, + is_numpy_namespace, + is_pydata_sparse_namespace, + is_torch_namespace, +) +from ._lib._helpers import eager_shape +from ._lib._typing import Array, ArrayNamespace + +__all__ = ["argpartition", "partition"] + + +def partition( + a: Array, + kth: int, + /, + axis: int | None = -1, + *, + xp: ArrayNamespace | None = None, +) -> Array: + """ + Return a partitioned copy of an array. + + Creates a copy of the array and partially sorts it in such a way that the value + of the element in k-th position is in the position it would be in a sorted array. + In the output array, all elements smaller than the k-th element are located to + the left of this element and all equal or greater are located to its right. + The ordering of the elements in the two partitions on the either side of + the k-th element in the output array is undefined. + + Parameters + ---------- + a : Array + Input array. + kth : int + Element index to partition by. + axis : int, optional + Axis along which to partition. The default is ``-1`` (the last axis). + If ``None``, the flattened array is used. + xp : array_namespace, optional + The standard-compatible namespace for `x`. Default: infer. + + Returns + ------- + partitioned_array + Array of the same type and shape as `a`. + + Notes + ----- + If `xp` implements ``partition`` or an equivalent function + (e.g. ``topk`` for torch), complexity will likely be O(n). + If not, this function simply calls ``xp.sort`` and complexity is O(n log n). + """ + # Validate inputs. + if xp is None: + xp = array_namespace(a) + if a.ndim < 1: + msg = "`a` must be at least 1-dimensional" + raise TypeError(msg) + if axis is None: + return partition(xp.reshape(a, (-1,)), kth, axis=0, xp=xp) + (size,) = eager_shape(a, axis) + if not (0 <= kth < size): + msg = f"kth(={kth}) out of bounds [0 {size})" + raise ValueError(msg) + + # Delegate where possible. + if is_numpy_namespace(xp) or is_cupy_namespace(xp) or is_jax_namespace(xp): + return xp.partition(a, kth, axis=axis) + + # Use top-k when possible: + if is_torch_namespace(xp): + if not (axis == -1 or axis == a.ndim - 1): + a = xp.transpose(a, axis, -1) + + out = xp.empty_like(a) + ranks = xp.arange(a.shape[-1]).expand_as(a) + + split_value, indices = xp.kthvalue(a, kth + 1, keepdim=True) + del indices # indices won't be used => del ASAP to reduce peak memory usage + + # fill the left-side of the partition + mask_src = a < split_value + n_left = mask_src.sum(dim=-1, keepdim=True) + mask_dest = ranks < n_left + out[mask_dest] = a[mask_src] + + # fill the middle of the partition + mask_src = a == split_value + n_left += mask_src.sum(dim=-1, keepdim=True) + mask_dest ^= ranks < n_left + out[mask_dest] = a[mask_src] + + # fill the right-side of the partition + mask_src = a > split_value + mask_dest = ranks >= n_left + out[mask_dest] = a[mask_src] + + if not (axis == -1 or axis == a.ndim - 1): + out = xp.transpose(out, axis, -1) + return out + + # Note: dask topk/argtopk sort the return values, so it's + # not much more efficient than sorting everything when + # kth is not small compared to x.size + + return _sorting.partition(a, kth, axis=axis, xp=xp) + + +def argpartition( + a: Array, + kth: int, + /, + axis: int | None = -1, + *, + xp: ArrayNamespace | None = None, +) -> Array: + """ + Perform an indirect partition along the given axis. + + It returns an array of indices of the same shape as `a` that + index data along the given axis in partitioned order. + + Parameters + ---------- + a : Array + Input array. + kth : int + Element index to partition by. + axis : int, optional + Axis along which to partition. The default is ``-1`` (the last axis). + If ``None``, the flattened array is used. + xp : array_namespace, optional + The standard-compatible namespace for `x`. Default: infer. + + Returns + ------- + index_array + Array of indices that partition `a` along the specified axis. + + Notes + ----- + If `xp` implements ``argpartition`` or an equivalent function + e.g. ``topk`` for torch), complexity will likely be O(n). + If not, this function simply calls ``xp.argsort`` and complexity is O(n log n). + """ + # Validate inputs. + if xp is None: + xp = array_namespace(a) + if is_pydata_sparse_namespace(xp): + msg = "Not implemented for sparse backend: no argsort" + raise NotImplementedError(msg) + if a.ndim < 1: + msg = "`a` must be at least 1-dimensional" + raise TypeError(msg) + if axis is None: + return argpartition(xp.reshape(a, (-1,)), kth, axis=0, xp=xp) + (size,) = eager_shape(a, axis) + if not (0 <= kth < size): + msg = f"kth(={kth}) out of bounds [0 {size})" + raise ValueError(msg) + + # Delegate where possible. + if is_numpy_namespace(xp) or is_cupy_namespace(xp) or is_jax_namespace(xp): + return xp.argpartition(a, kth, axis=axis) + + # Use top-k when possible: + if is_torch_namespace(xp): + # see `partition` above for commented details of those steps: + if not (axis == -1 or axis == a.ndim - 1): + a = xp.transpose(a, axis, -1) + + ranks = xp.arange(a.shape[-1]).expand_as(a) + out = xp.empty_like(ranks) + + split_value, indices = xp.kthvalue(a, kth + 1, keepdim=True) + del indices # indices won't be used => del ASAP to reduce peak memory usage + + mask_src = a < split_value + n_left = mask_src.sum(dim=-1, keepdim=True) + mask_dest = ranks < n_left + out[mask_dest] = ranks[mask_src] + + mask_src = a == split_value + n_left += mask_src.sum(dim=-1, keepdim=True) + mask_dest ^= ranks < n_left + out[mask_dest] = ranks[mask_src] + + mask_src = a > split_value + mask_dest = ranks >= n_left + out[mask_dest] = ranks[mask_src] + + if not (axis == -1 or axis == a.ndim - 1): + out = xp.transpose(out, axis, -1) + return out + + # Note: dask topk/argtopk sort the return values, so it's + # not much more efficient than sorting everything when + # kth is not small compared to x.size + + return _sorting.argpartition(a, kth, axis=axis, xp=xp) diff --git a/src/array_api_extra/_statistical.py b/src/array_api_extra/_statistical.py new file mode 100644 index 00000000..01c2b3ac --- /dev/null +++ b/src/array_api_extra/_statistical.py @@ -0,0 +1,380 @@ +"""Delegation layer for statistical functions.""" + +from ._agnostic import _statistical +from ._lib._compat import ( + array_namespace, + is_cupy_namespace, + is_dask_namespace, + is_jax_namespace, + is_numpy_namespace, + is_torch_namespace, +) +from ._lib._typing import Array, ArrayNamespace + +__all__ = ["cov", "nanmax", "nanmin", "nansum"] + + +def cov( + m: Array, + /, + *, + axis: int = -1, + correction: float = 1, + fweights: Array | None = None, + aweights: Array | None = None, + xp: ArrayNamespace | None = None, +) -> Array: + """ + Estimate a covariance matrix (or a stack of covariance matrices). + + Covariance indicates the level to which two variables vary together. + If we examine *N*-dimensional samples, :math:`X = [x_1, x_2, ... x_N]^T`, + each with *M* observations, then element :math:`C_{ij}` of the + :math:`N \\times N` covariance matrix is the covariance of + :math:`x_i` and :math:`x_j`. The element :math:`C_{ii}` is the variance + of :math:`x_i`. + + Extends :func:`numpy.cov` with support for batch input. + Naming follows the array API conventions used elsewhere in + this library (``axis``, ``correction``) rather than the NumPy spellings + (``rowvar``, ``bias``, ``ddof``); see Notes for the mapping. + + Parameters + ---------- + m : array + An array of shape ``(..., N, M)`` whose innermost two dimensions + contain *M* observations of *N* variables by default. The axis of + observations is controlled by `axis`. + axis : int, optional + Axis of `m` containing the observations. Default: ``-1`` (the last + axis), matching the array API convention. Use ``axis=-2`` (or ``0`` + for 2-D input) to treat each column as a variable, which + corresponds to ``rowvar=False`` in :func:`numpy.cov`. + correction : int or float, optional + Degrees of freedom correction: normalization divides by + ``N - correction`` (for unweighted input). Default: ``1``, which + gives the unbiased estimate (matches :func:`numpy.cov` default of + ``bias=False``). Set to ``0`` for the biased estimate (``N`` + normalization). Corresponds to ``ddof`` in :func:`numpy.cov` and to + ``correction`` in :func:`numpy.var`/:func:`numpy.std` and + :func:`torch.cov`. + Non-integer values are allowed for advanced use cases: the + unbiased correction for weighted observations depends on the + sum and dispersion of the weights and is generally not an + integer, and autocorrelated data may also require a fractional + correction. Non-integer ``correction`` routes through the + generic implementation because :func:`numpy.cov`'s ``ddof`` and + :func:`torch.cov`'s ``correction`` both require integers. + fweights : array, optional + 1-D array of integer frequency weights: the number of times each + observation is repeated. Same as ``fweights`` in + :func:`numpy.cov`/:func:`torch.cov`. + aweights : array, optional + 1-D array of observation-vector weights (analytic weights). Larger + values mark more important observations. Same as ``aweights`` in + :func:`numpy.cov`/:func:`torch.cov`. + xp : array_namespace, optional + The standard-compatible namespace for `m`. Default: infer. + + Returns + ------- + array + An array having shape ``(..., N, N)`` whose innermost two dimensions represent + the covariance matrix of the variables. + + Notes + ----- + Mapping from :func:`numpy.cov` to this function:: + + numpy.cov(m, rowvar=True) -> cov(m, axis=-1) # default + numpy.cov(m, rowvar=False) -> cov(m, axis=-2) + numpy.cov(m, bias=True) -> cov(m, correction=0) + numpy.cov(m, ddof=k) -> cov(m, correction=k) + numpy.cov(m, fweights=f) -> cov(m, fweights=f) + numpy.cov(m, aweights=a) -> cov(m, aweights=a) + + A ``RuntimeWarning`` is emitted for non-positive effective degrees of + freedom when the effective normalizer can be checked without materializing + a lazy array. When the normalizer itself is lazy (e.g. for weighted Dask + inputs), this check is skipped; choose ``correction`` and weights such that + it is positive. + + Examples + -------- + >>> import array_api_strict as xp + >>> import array_api_extra as xpx + + Consider two variables, :math:`x_0` and :math:`x_1`, which + correlate perfectly, but in opposite directions: + + >>> x = xp.asarray([[0, 2], [1, 1], [2, 0]]).T + >>> x + Array([[0, 1, 2], + [2, 1, 0]], dtype=array_api_strict.int64) + + Note how :math:`x_0` increases while :math:`x_1` decreases. The covariance + matrix shows this clearly: + + >>> xpx.cov(x, xp=xp) + Array([[ 1., -1.], + [-1., 1.]], dtype=array_api_strict.float64) + + Note that element :math:`C_{0,1}`, which shows the correlation between + :math:`x_0` and :math:`x_1`, is negative. + + Further, note how `x` and `y` are combined: + + >>> x = xp.asarray([-2.1, -1, 4.3]) + >>> y = xp.asarray([3, 1.1, 0.12]) + >>> X = xp.stack((x, y), axis=0) + >>> xpx.cov(X, xp=xp) + Array([[11.71 , -4.286 ], + [-4.286 , 2.14413333]], dtype=array_api_strict.float64) + + >>> xpx.cov(x, xp=xp) + Array(11.71, dtype=array_api_strict.float64) + + >>> xpx.cov(y, xp=xp) + Array(2.14413333, dtype=array_api_strict.float64) + + Input with more than two dimensions is treated as a stack of + two-dimensional input. + + >>> stack = xp.stack((X, 2*X)) + >>> xpx.cov(stack) + Array([[[ 11.71 , -4.286 ], + [ -4.286 , 2.14413333]], + [[ 46.84 , -17.144 ], + [-17.144 , 8.57653333]]], dtype=array_api_strict.float64) + + The normalization can be adjusted with `correction`, and observations + can be weighted with integer frequencies `fweights` or importance + weights `aweights`: + + >>> x = xp.asarray([0., 1., 2., 3., 4.]) + >>> xpx.cov(x, xp=xp) # unbiased variance: divide by N - 1 + Array(2.5, dtype=array_api_strict.float64) + >>> xpx.cov(x, correction=0, xp=xp) # biased variance: divide by N + Array(2., dtype=array_api_strict.float64) + + Giving the two extreme observations frequency 2 via `fweights` is + equivalent to repeating them in `x`: + + >>> xpx.cov(x, fweights=xp.asarray([2, 1, 1, 1, 2]), xp=xp) + Array(3., dtype=array_api_strict.float64) + >>> xpx.cov(xp.asarray([0., 0., 1., 2., 3., 4., 4.]), xp=xp) + Array(3., dtype=array_api_strict.float64) + + `aweights` instead adjusts the relative importance of observations, + here down-weighting the two extremes: + + >>> xpx.cov(x, aweights=xp.asarray([0.5, 1., 1., 1., 0.5]), xp=xp) + Array(1.92, dtype=array_api_strict.float64) + """ + + if xp is None: + xp = array_namespace(m, fweights, aweights) + + # Validate axis against m.ndim. + ndim = max(m.ndim, 1) + if not -ndim <= axis < ndim: + msg = f"axis {axis} is out of bounds for array of dimension {m.ndim}" + raise IndexError(msg) + + # Normalize: observations on the last axis. After this, every backend + # sees the same convention and we never need to deal with `rowvar`. + if m.ndim >= 2 and axis not in (-1, m.ndim - 1): + m = xp.moveaxis(m, axis, -1) + + # `numpy.cov` (and cupy/dask/jax) require integer `ddof`; `torch.cov` + # requires integer `correction`. For non-integer-valued `correction`, + # fall through to the generic implementation. + integer_correction = float(correction).is_integer() + has_weights = fweights is not None or aweights is not None + + if m.ndim <= 2 and integer_correction: + # Not just for static typing: `correction` may be an integer-valued + # float such as 1.0, which `torch.cov` rejects at runtime. + int_correction = int(correction) + if is_torch_namespace(xp): + fw = None if fweights is None else xp.asarray(fweights) + aw = None if aweights is None else xp.asarray(aweights) + return xp.cov(m, correction=int_correction, fweights=fw, aweights=aw) + # `dask.array.cov` forces `.compute()` whenever weights are given: + # its internal `if fact <= 0` check on a lazy 0-D scalar triggers + # materialization. Route to the generic impl, which is fully lazy + # because it only does sum/matmul and skips that scalar check. + if ( + is_numpy_namespace(xp) + or is_cupy_namespace(xp) + or is_jax_namespace(xp) + or (is_dask_namespace(xp) and not has_weights) + ): + return xp.cov( + m, + ddof=int_correction, + fweights=fweights, + aweights=aweights, + ) + + return _statistical.cov( + m, + correction=correction, + fweights=fweights, + aweights=aweights, + xp=xp, + ) + + +def nanmin( + a: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + xp: ArrayNamespace | None = None, +) -> Array: + """ + Return the minimum of the array elements along a given axis, ignoring NaNs. + + Parameters + ---------- + a : Array + Input array. + axis : int or tuple of ints or None, optional + Axis or axes along which the minimum is computed. The default is to compute + the minimum of the flattened array. + xp : array_namespace, optional + The standard-compatible namespace for `a`. Default: infer. + + Returns + ------- + array + An array of minimum values along the given axis, ignoring NaNs. + + Examples + -------- + >>> import array_api_extra as xpx + >>> import array_api_strict as xp + >>> a = xp.asarray([[5, 3, xp.nan, 1], [4, xp.nan, 2, xp.nan]]) + >>> xpx.nanmin(a) + Array(1., dtype=array_api_strict.float64) + >>> xpx.nanmin(a, axis=0) + Array([4., 3., 2., 1.], dtype=array_api_strict.float64) + >>> xpx.nanmin(a, axis=1) + Array([1., 2.], dtype=array_api_strict.float64) + """ + if xp is None: + xp = array_namespace(a) + + if ( + is_numpy_namespace(xp) + or is_cupy_namespace(xp) + or is_dask_namespace(xp) + or is_jax_namespace(xp) + ): + return xp.nanmin(a, axis=axis) + + return _statistical.nanmin(a, axis=axis, xp=xp) + + +def nanmax( + a: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + xp: ArrayNamespace | None = None, +) -> Array: + """ + Return the maximum of the array elements along a given axis, ignoring NaNs. + + Parameters + ---------- + a : Array + Input array. + axis : int or tuple of ints or None, optional + Axis or axes along which the maximum is computed. The default is to compute + the maximum of the flattened array. + xp : array_namespace, optional + The standard-compatible namespace for `a`. Default: infer. + + Returns + ------- + array + An array of maximum values along the given axis, ignoring NaNs. + + Examples + -------- + >>> import array_api_extra as xpx + >>> import array_api_strict as xp + >>> a = xp.asarray([[5, 3, xp.nan, 6], [4, xp.nan, 2, xp.nan]]) + >>> xpx.nanmax(a) + Array(6., dtype=array_api_strict.float64) + >>> xpx.nanmax(a, axis=0) + Array([5., 3., 2., 6.], dtype=array_api_strict.float64) + >>> xpx.nanmax(a, axis=1) + Array([6., 4.], dtype=array_api_strict.float64) + """ + if xp is None: + xp = array_namespace(a) + + if ( + is_numpy_namespace(xp) + or is_cupy_namespace(xp) + or is_dask_namespace(xp) + or is_jax_namespace(xp) + ): + return xp.nanmax(a, axis=axis) + + return _statistical.nanmax(a, axis=axis, xp=xp) + + +def nansum( + a: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + xp: ArrayNamespace | None = None, +) -> Array: + """ + Return the sum of the array elements along a given axis, ignoring NaNs. + + Parameters + ---------- + a : Array + Input array. + axis : int or tuple of ints or None, optional + Axis or axes along which the sum is computed. The default is to compute + the sum of the flattened array. + xp : array_namespace, optional + The standard-compatible namespace for `a`. Default: infer. + + Returns + ------- + array + An array of sum values along the given axis, ignoring NaNs. + + Examples + -------- + >>> import array_api_extra as xpx + >>> import array_api_strict as xp + >>> a = xp.asarray([[5, 3, xp.nan, 1], [4, xp.nan, 2, xp.nan]]) + >>> xpx.nansum(a) + Array(15., dtype=array_api_strict.float64) + >>> xpx.nansum(a, axis=0) + Array([9., 3., 2., 1.], dtype=array_api_strict.float64) + >>> xpx.nansum(a, axis=1) + Array([9., 6.], dtype=array_api_strict.float64) + """ + if xp is None: + xp = array_namespace(a) + + if ( + is_numpy_namespace(xp) + or is_cupy_namespace(xp) + or is_dask_namespace(xp) + or is_jax_namespace(xp) + or is_torch_namespace(xp) + ): + return xp.nansum(a, axis=axis) + + return _statistical.nansum(a, axis=axis, xp=xp) diff --git a/src/array_api_extra/testing/__init__.py b/src/array_api_extra/testing/__init__.py new file mode 100644 index 00000000..a57a88a6 --- /dev/null +++ b/src/array_api_extra/testing/__init__.py @@ -0,0 +1,19 @@ +"""Public testing utilities.""" + +from ._testing import ( + assert_close, + assert_close_nulp, + assert_equal, + assert_less, + lazy_xp_function, + patch_lazy_xp_functions, +) + +__all__ = [ + "assert_close", + "assert_close_nulp", + "assert_equal", + "assert_less", + "lazy_xp_function", + "patch_lazy_xp_functions", +] diff --git a/src/array_api_extra/testing.py b/src/array_api_extra/testing/_testing.py similarity index 98% rename from src/array_api_extra/testing.py rename to src/array_api_extra/testing/_testing.py index 259b8831..79e9ccf3 100644 --- a/src/array_api_extra/testing.py +++ b/src/array_api_extra/testing/_testing.py @@ -1,8 +1,4 @@ -""" -Public testing utilities. - -See also _lib._testing for additional private testing utilities. -""" +"""Implementation of `array_api_extra.testing`.""" from __future__ import annotations @@ -16,7 +12,7 @@ from types import FunctionType, ModuleType from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar, cast -from ._lib._utils._compat import ( +from .._lib._compat import ( array_namespace, is_array_api_strict_namespace, is_cupy_namespace, @@ -28,11 +24,11 @@ is_torch_namespace, to_device, ) -from ._lib._utils._compat import ( +from .._lib._compat import ( device as get_device, ) -from ._lib._utils._helpers import jax_autojit, pickle_flatten, pickle_unflatten -from ._lib._utils._typing import Array, ArrayNamespace, Device +from .._lib._helpers import jax_autojit, pickle_flatten, pickle_unflatten +from .._lib._typing import Array, ArrayNamespace, Device __all__ = [ "assert_close", @@ -48,29 +44,39 @@ import numpy as np import pytest from dask.typing import Graph, Key, SchedulerGetCallable - from typing_extensions import override + from typing_extensions import override as _override else: # Sphinx hacks SchedulerGetCallable = object - def override(func): + def _override(func): return func +__all__ = [ + "assert_close", + "assert_close_nulp", + "assert_equal", + "assert_less", + "lazy_xp_function", + "patch_lazy_xp_functions", +] + + P = ParamSpec("P") T = TypeVar("T") _ufuncs_tags: dict[object, dict[str, Any]] = {} -class Deprecated(enum.Enum): +class _Deprecated(enum.Enum): """Unique type for deprecated parameters.""" DEPRECATED = 1 -DEPRECATED = Deprecated.DEPRECATED +DEPRECATED = _Deprecated.DEPRECATED def _clone_function( # numpydoc ignore=PR01,RT01 @@ -93,8 +99,8 @@ def lazy_xp_function( *, allow_dask_compute: bool | int = False, jax_jit: bool = True, - static_argnums: Deprecated = DEPRECATED, - static_argnames: Deprecated = DEPRECATED, + static_argnums: _Deprecated = DEPRECATED, + static_argnames: _Deprecated = DEPRECATED, ) -> None: # numpydoc ignore=GL07 """ Tag a function to be tested on lazy backends. @@ -501,7 +507,7 @@ def __init__(self, max_count: int, msg: str) -> None: # numpydoc ignore=GL08 self.max_count = max_count self.msg = msg - @override + @_override def __call__( self, dsk: Graph, keys: Sequence[Key] | Key, **kwargs: Any ) -> Any: # numpydoc ignore=GL08 diff --git a/tests/conftest.py b/tests/conftest.py index 49b5960e..7a83fe37 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -8,10 +8,10 @@ import pytest from array_api_extra._lib._backends import Backend +from array_api_extra._lib._compat import array_namespace +from array_api_extra._lib._compat import device as get_device from array_api_extra._lib._testing import xfail -from array_api_extra._lib._utils._compat import array_namespace -from array_api_extra._lib._utils._compat import device as get_device -from array_api_extra._lib._utils._typing import ArrayNamespace, Device +from array_api_extra._lib._typing import ArrayNamespace, Device from array_api_extra.testing import patch_lazy_xp_functions T = TypeVar("T") diff --git a/tests/meson.build b/tests/meson.build index 246128f1..911502a4 100644 --- a/tests/meson.build +++ b/tests/meson.build @@ -2,10 +2,20 @@ py.install_sources([ '__init__.py', 'conftest.py', 'test_at.py', + 'test_creation.py', 'test_deprecation.py', - 'test_funcs.py', + 'test_elementwise.py', 'test_helpers.py', + 'test_indexing.py', + 'test_inspection.py', 'test_lazy.py', + 'test_linalg.py', + 'test_manipulation.py', + 'test_public_api.py', + 'test_searching.py', + 'test_set.py', + 'test_sorting.py', + 'test_statistical.py', 'test_testing.py', 'test_version.py', ], diff --git a/tests/test_at.py b/tests/test_at.py index 987eddbc..24f52002 100644 --- a/tests/test_at.py +++ b/tests/test_at.py @@ -7,11 +7,11 @@ import pytest from array_api_extra import at -from array_api_extra._lib._at import _AtOp +from array_api_extra._at import _AtOp from array_api_extra._lib._backends import Backend -from array_api_extra._lib._utils._compat import array_namespace, is_writeable_array -from array_api_extra._lib._utils._compat import device as get_device -from array_api_extra._lib._utils._typing import ( +from array_api_extra._lib._compat import array_namespace, is_writeable_array +from array_api_extra._lib._compat import device as get_device +from array_api_extra._lib._typing import ( Array, ArrayNamespace, Device, diff --git a/tests/test_creation.py b/tests/test_creation.py new file mode 100644 index 00000000..358efa09 --- /dev/null +++ b/tests/test_creation.py @@ -0,0 +1,174 @@ +import numpy as np +import pytest + +from array_api_extra import at, create_diagonal, one_hot +from array_api_extra._lib._backends import Backend +from array_api_extra._lib._compat import device as get_device +from array_api_extra._lib._helpers import eager_shape, ndindex +from array_api_extra._lib._typing import ArrayNamespace, Device +from array_api_extra.testing import assert_equal, lazy_xp_function + +lazy_xp_function(create_diagonal) +lazy_xp_function(one_hot) + + +@pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no arange", strict=False) +class TestOneHot: + @pytest.mark.parametrize("n_dim", range(4)) + @pytest.mark.parametrize("num_classes", [1, 3, 10]) + def test_dims_and_classes(self, xp: ArrayNamespace, n_dim: int, num_classes: int): + shape = tuple(range(2, 2 + n_dim)) + rng = np.random.default_rng(2347823) + np_x = rng.integers(num_classes, size=shape) + x = xp.asarray(np_x) + y = one_hot(x, num_classes) + assert y.shape == (*x.shape, num_classes) + for *i_list, j in ndindex(*shape, num_classes): + i = tuple(i_list) + assert float(y[(*i, j)]) == (int(x[i]) == j) + + def test_basic(self, xp: ArrayNamespace): + actual = one_hot(xp.asarray([0, 1, 2]), 3) + expected = xp.asarray([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]) + assert_equal(actual, expected) + + actual = one_hot(xp.asarray([1, 2, 0]), 3) + expected = xp.asarray([[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]]) + assert_equal(actual, expected) + + def test_2d(self, xp: ArrayNamespace): + actual = one_hot(xp.asarray([[2, 1, 0], [1, 0, 2]]), 3, axis=1) + expected = xp.asarray( + [ + [[0.0, 0.0, 1.0], [0.0, 1.0, 0.0], [1.0, 0.0, 0.0]], + [[0.0, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]], + ] + ) + assert_equal(actual, expected) + + @pytest.mark.skip_xp_backend( + Backend.ARRAY_API_STRICTEST, reason="backend doesn't support Boolean indexing" + ) + def test_abstract_size(self, xp: ArrayNamespace): + x = xp.arange(5) + x = x[x > 2] + actual = one_hot(x, 5) + expected = xp.asarray([[0.0, 0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 0.0, 1.0]]) + assert_equal(actual, expected) + + @pytest.mark.skip_xp_backend( + Backend.TORCH_GPU, reason="Puts Pytorch into a bad state." + ) + def test_out_of_bound(self, xp: ArrayNamespace): + # Undefined behavior. Either return zero, or raise. + try: + actual = one_hot(xp.asarray([-1, 3]), 3) + except IndexError: + return + expected = xp.asarray([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]) + assert_equal(actual, expected) + + @pytest.mark.parametrize( + "int_dtype", + ["int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64"], + ) + def test_int_types(self, xp: ArrayNamespace, int_dtype: str): + dtype = getattr(xp, int_dtype) + x = xp.asarray([0, 1, 2], dtype=dtype) + actual = one_hot(x, 3) + expected = xp.asarray([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]) + assert_equal(actual, expected) + + def test_custom_dtype(self, xp: ArrayNamespace): + actual = one_hot(xp.asarray([0, 1, 2], dtype=xp.int32), 3, dtype=xp.bool) + expected = xp.asarray( + [[True, False, False], [False, True, False], [False, False, True]] + ) + assert_equal(actual, expected) + + def test_axis(self, xp: ArrayNamespace): + expected = xp.asarray([[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]]).T + actual = one_hot(xp.asarray([1, 2, 0]), 3, axis=0) + assert_equal(actual, expected) + + actual = one_hot(xp.asarray([1, 2, 0]), 3, axis=-2) + assert_equal(actual, expected) + + def test_non_integer(self, xp: ArrayNamespace): + with pytest.raises(TypeError): + _ = one_hot(xp.asarray([1.0]), 3) + + def test_device(self, xp: ArrayNamespace, device: Device): + x = xp.asarray([0, 1, 2], device=device) + y = one_hot(x, 3) + assert get_device(y) == device + + +@pytest.mark.skip_xp_backend( + Backend.SPARSE, reason="read-only backend without .at support" +) +class TestCreateDiagonal: + def test_1d_from_numpy(self, xp: ArrayNamespace): + # from np.diag tests + vals = 100 * xp.arange(5, dtype=xp.float64) + b = xp.zeros((5, 5), dtype=xp.float64) + for k in range(5): + b = at(b)[k, k].set(vals[k]) + assert_equal(create_diagonal(vals), b) + b = xp.zeros((7, 7), dtype=xp.float64) + c = xp.asarray(b, copy=True) + for k in range(5): + b = at(b)[k, k + 2].set(vals[k]) + c = at(c)[k + 2, k].set(vals[k]) + assert_equal(create_diagonal(vals, offset=2), b) + assert_equal(create_diagonal(vals, offset=-2), c) + + @pytest.mark.parametrize("n", range(1, 10)) + @pytest.mark.parametrize("offset", range(1, 10)) + def test_1d_from_scipy(self, xp: ArrayNamespace, n: int, offset: int): + # from scipy._lib tests + rng = np.random.default_rng(2347823) + one = xp.asarray(1.0) + x = rng.random(n) + A = create_diagonal(xp.asarray(x, dtype=one.dtype), offset=offset) + B = xp.asarray(np.diag(x, offset), dtype=one.dtype) + assert_equal(A, B) + + def test_0d_raises(self, xp: ArrayNamespace): + with pytest.raises(ValueError, match="1-dimensional"): + _ = create_diagonal(xp.asarray(1)) + + @pytest.mark.parametrize( + "shape", + [ + (0,), + (10,), + (0, 1), + (1, 0), + (0, 0), + (2, 3), + (4, 2, 1), + (1, 1, 7), + (0, 0, 1), + (3, 2, 4, 5), + ], + ) + def test_nd(self, xp: ArrayNamespace, shape: tuple[int, ...]): + rng = np.random.default_rng(2347823) + b = xp.asarray( + rng.integers((1 << 64) - 1, size=shape, dtype=np.uint64), dtype=xp.uint64 + ) + c = create_diagonal(b) + zero = xp.zeros((), dtype=xp.uint64) + assert c.shape == (*b.shape, b.shape[-1]) + for i in ndindex(*eager_shape(c)): + assert_equal(c[i], b[i[:-1]] if i[-2] == i[-1] else zero) + + def test_device(self, xp: ArrayNamespace, device: Device): + x = xp.asarray([1, 2, 3], device=device) + assert get_device(create_diagonal(x)) == device + + def test_xp(self, xp: ArrayNamespace): + x = xp.asarray([1, 2]) + y = create_diagonal(x, xp=xp) + assert_equal(y, xp.asarray([[1, 0], [0, 2]])) diff --git a/tests/test_deprecation.py b/tests/test_deprecation.py index 1fd4c310..5425cf58 100644 --- a/tests/test_deprecation.py +++ b/tests/test_deprecation.py @@ -1,7 +1,7 @@ import pytest from array_api_extra import broadcast_shapes, expand_dims -from array_api_extra._lib._utils._typing import ArrayNamespace +from array_api_extra._lib._typing import ArrayNamespace class TestDeprecatedFunctions: diff --git a/tests/test_elementwise.py b/tests/test_elementwise.py new file mode 100644 index 00000000..43a2e530 --- /dev/null +++ b/tests/test_elementwise.py @@ -0,0 +1,762 @@ +import warnings +from typing import Any, cast + +import hypothesis +import hypothesis.extra.numpy as npst +import numpy as np +import pytest +from hypothesis import given +from hypothesis import strategies as st + +from array_api_extra import ( + angle, + apply_where, + default_dtype, + deg2rad, + isclose, + nan_to_num, + rad2deg, + sinc, +) +from array_api_extra._lib._backends import NUMPY_VERSION, Backend +from array_api_extra._lib._compat import device as get_device +from array_api_extra._lib._typing import Array, ArrayNamespace, Device +from array_api_extra.testing import assert_close, assert_equal, lazy_xp_function + +lazy_xp_function(apply_where) +lazy_xp_function(deg2rad) +lazy_xp_function(isclose) +lazy_xp_function(nan_to_num) +lazy_xp_function(rad2deg) +lazy_xp_function(sinc) + + +class TestApplyWhere: + @staticmethod + def f1(x: Array, y: Array | int = 10) -> Array: + return x + y + + @staticmethod + def f2(x: Array, y: Array | int = 10) -> Array: + return x - y + + def test_f1_f2(self, xp: ArrayNamespace): + x = xp.asarray([1, 2, 3, 4]) + cond = x % 2 == 0 + actual = apply_where(cond, x, self.f1, self.f2) + expect = xp.where(cond, self.f1(x), self.f2(x)) + assert_equal(actual, expect) + + def test_fill_value(self, xp: ArrayNamespace): + x = xp.asarray([1, 2, 3, 4]) + cond = x % 2 == 0 + actual = apply_where(x % 2 == 0, x, self.f1, fill_value=0) + expect = xp.where(cond, self.f1(x), xp.asarray(0)) + assert_equal(actual, expect) + + actual = apply_where(x % 2 == 0, x, self.f1, fill_value=xp.asarray(0)) + assert_equal(actual, expect) + + def test_args_tuple(self, xp: ArrayNamespace): + x = xp.asarray([1, 2, 3, 4]) + y = xp.asarray([10, 20, 30, 40]) + cond = x % 2 == 0 + actual = apply_where(cond, (x, y), self.f1, self.f2) + expect = xp.where(cond, self.f1(x, y), self.f2(x, y)) + assert_equal(actual, expect) + + def test_broadcast(self, xp: ArrayNamespace): + x = xp.asarray([1, 2]) + y = xp.asarray([[10], [20], [30]]) + cond = xp.broadcast_to(xp.asarray(True), (4, 1, 1)) + + actual = apply_where(cond, (x, y), self.f1, self.f2) + expect = xp.where(cond, self.f1(x, y), self.f2(x, y)) + assert_equal(actual, expect) + + actual = apply_where( + cond, + (x, y), + lambda x, _: x, + lambda _, y: y, + ) + expect = xp.where(cond, x, y) + assert_equal(actual, expect) + + # Shaped fill_value + actual = apply_where(cond, x, self.f1, fill_value=y) + expect = xp.where(cond, self.f1(x), y) + assert_equal(actual, expect) + + def test_dtype_propagation(self, xp: ArrayNamespace, library: Backend): + x = xp.asarray([1, 2], dtype=xp.int8) + y = xp.asarray([3, 4], dtype=xp.int16) + cond = x % 2 == 0 + + mxp = np if library is Backend.DASK else xp + actual = apply_where( + cond, + (x, y), + self.f1, + lambda x, y: mxp.astype(x - y, xp.int64), # pyright: ignore[reportArgumentType] # pyrefly: ignore[bad-argument-type] + ) + assert actual.dtype == xp.int64 + + actual = apply_where(cond, y, self.f1, fill_value=5) + assert actual.dtype == xp.int16 + + @pytest.mark.parametrize("fill_value_raw", [3, [3, 4]]) + @pytest.mark.parametrize( + ("fill_value_dtype", "expect_dtype"), [("int32", "int32"), ("int8", "int16")] + ) + def test_dtype_propagation_fill_value( + self, + xp: ArrayNamespace, + fill_value_raw: int | list[int], + fill_value_dtype: str, + expect_dtype: str, + ): + x = xp.asarray([1, 2], dtype=xp.int16) + cond = x % 2 == 0 + fill_value = xp.asarray(fill_value_raw, dtype=getattr(xp, fill_value_dtype)) + + actual = apply_where(cond, x, self.f1, fill_value=fill_value) + assert actual.dtype == getattr(xp, expect_dtype) + + def test_dont_overwrite_fill_value(self, xp: ArrayNamespace): + x = xp.asarray([1, 2]) + fill_value = xp.asarray([100, 200]) + actual = apply_where(x % 2 == 0, x, self.f1, fill_value=fill_value) + assert_equal(actual, xp.asarray([100, 12])) + assert_equal(fill_value, xp.asarray([100, 200])) + + @pytest.mark.skip_xp_backend( + Backend.ARRAY_API_STRICTEST, + reason="no boolean indexing -> run everywhere", + ) + @pytest.mark.skip_xp_backend( + Backend.SPARSE, + reason="no indexing by sparse array -> run everywhere", + ) + def test_dont_run_on_false(self, xp: ArrayNamespace): + x = xp.asarray([1.0, 2.0, 0.0]) + y = xp.asarray([0.0, 3.0, 4.0]) + # On NumPy, division by zero will trigger warnings + actual = apply_where( + x == 0, + (x, y), + lambda x, y: x / y, + lambda x, y: y / x, + ) + assert_equal(actual, xp.asarray([0.0, 1.5, 0.0])) + + def test_bad_args(self, xp: ArrayNamespace): + x = xp.asarray([1, 2, 3, 4]) + cond = x % 2 == 0 + # Neither f2 nor fill_value + with pytest.raises(TypeError, match="Exactly one of"): + apply_where(cond, x, self.f1) # type: ignore[call-overload] # pyright: ignore[reportCallIssue] + # Both f2 and fill_value + with pytest.raises(TypeError, match="Exactly one of"): + apply_where(cond, x, self.f1, self.f2, fill_value=0) # type: ignore[call-overload] # pyright: ignore[reportCallIssue] + + @pytest.mark.skip_xp_backend(Backend.NUMPY_READONLY, reason="xp=xp") + def test_xp(self, xp: ArrayNamespace): + x = xp.asarray([1, 2, 3, 4]) + cond = x % 2 == 0 + actual = apply_where(cond, x, self.f1, self.f2, xp=xp) + expect = xp.where(cond, self.f1(x), self.f2(x)) + assert_equal(actual, expect) + + def test_device(self, xp: ArrayNamespace, device: Device): + x = xp.asarray([1, 2, 3, 4], device=device) + y = apply_where(x % 2 == 0, x, self.f1, self.f2) + assert get_device(y) == device + y = apply_where(x % 2 == 0, x, self.f1, fill_value=0) + assert get_device(y) == device + y = apply_where(x % 2 == 0, x, self.f1, fill_value=x) + assert get_device(y) == device + + @pytest.mark.filterwarnings("ignore::RuntimeWarning") # overflows, etc. + @hypothesis.settings( + # The xp and library fixtures are not regenerated between hypothesis iterations + suppress_health_check=[hypothesis.HealthCheck.function_scoped_fixture], + # JAX can take a long time to initialize on the first call + deadline=None, + ) + @given( + n_arrays=st.integers(min_value=0, max_value=3), + n_kwarrays=st.integers(min_value=0, max_value=3), + rng_seed=st.integers(min_value=1000000000, max_value=9999999999), + dtype=npst.floating_dtypes(sizes=(32, 64)), + p=st.floats(min_value=0, max_value=1), + data=st.data(), + ) + def test_hypothesis( + self, + n_arrays: int, + n_kwarrays: int, + rng_seed: int, + dtype: np.dtype[Any], + p: float, + data: st.DataObject, + xp: ArrayNamespace, + library: Backend, + ): + if ( + library.like(Backend.NUMPY) + and NUMPY_VERSION < (2, 0) + and dtype.type is np.float32 + ): + pytest.xfail(reason="NumPy 1.x dtype promotion for scalars") + + _ = hypothesis.assume(n_arrays + n_kwarrays > 0) + mbs = npst.mutually_broadcastable_shapes( + num_shapes=1 + n_arrays + n_kwarrays, min_side=0 + ) + input_shapes, _ = data.draw(mbs) + cond_shape = input_shapes[0] + shapes = input_shapes[1 : 1 + n_arrays] + kwshapes = input_shapes[1 + n_arrays :] + + # cupy/cupy#8382 + # https://github.com/jax-ml/jax/issues/26658 + elements = {"allow_subnormal": not library.like(Backend.CUPY, Backend.JAX)} + + fill_value = xp.asarray( + data.draw(npst.arrays(dtype=dtype.type, shape=(), elements=elements)) + ) + float_fill_value = float(fill_value) + if library is Backend.CUPY and dtype.type is np.float32: + # Avoid data-dependent dtype promotion when encountering subnormals + # close to the max float32 value + float_fill_value = float(np.clip(float_fill_value, -1e38, 1e38)) + + arrays = tuple( + xp.asarray( + data.draw(npst.arrays(dtype=dtype.type, shape=shape, elements=elements)) + ) + for shape in shapes + ) + + kwargs = { + f"kw{n}": xp.asarray( + data.draw(npst.arrays(dtype=dtype.type, shape=shape, elements=elements)) + ) + for n, shape in enumerate(kwshapes) + } + kwkeys = kwargs.keys() + + def f1(*args: Array, **kwargs: dict[str, Array]) -> Array: + assert kwargs.keys() == kwkeys + args_kwargs = cast(tuple[Array, ...], (*args, *kwargs.values())) + return cast(Array, sum(args_kwargs)) + + def f2(*args: Array, **kwargs: dict[str, Array]) -> Array: + assert kwargs.keys() == kwkeys + args_kwargs = cast(tuple[Array, ...], (*args, *kwargs.values())) + return cast(Array, sum(args_kwargs) / 2) + + rng = np.random.default_rng(rng_seed) + cond = xp.asarray(rng.random(size=cond_shape) > p) + + res1 = apply_where(cond, arrays, f1, fill_value=fill_value, kwargs=kwargs) + res2 = apply_where(cond, arrays, f1, f2, kwargs=kwargs) + res3 = apply_where(cond, arrays, f1, fill_value=float_fill_value, kwargs=kwargs) + + ref1 = xp.where(cond, f1(*arrays, **kwargs), fill_value) + ref2 = xp.where(cond, f1(*arrays, **kwargs), f2(*arrays, **kwargs)) + ref3 = xp.where(cond, f1(*arrays, **kwargs), float_fill_value) + + assert_close(res1, ref1, rtol=2e-16) + assert_equal(res2, ref2) + assert_equal(res3, ref3) + + +@pytest.mark.filterwarnings( # array_api_strictest + "ignore:invalid value encountered:RuntimeWarning:array_api_strict" +) +@pytest.mark.filterwarnings( # sparse + "ignore:invalid value encountered:RuntimeWarning:sparse" +) +class TestIsClose: + @pytest.mark.parametrize("swap", [False, True]) + @pytest.mark.parametrize( + ("a", "b"), + [ + (0.0, 0.0), + (1.0, 1.0), + (1.0, 2.0), + (1.0, -1.0), + (100.0, 101.0), + (0, 0), + (1, 1), + (1, 2), + (1, -1), + (1.0 + 1j, 1.0 + 1j), + (1.0 + 1j, 1.0 - 1j), + (float("inf"), float("inf")), + (float("inf"), 100.0), + (float("inf"), float("-inf")), + (float("-inf"), float("-inf")), + (float("nan"), float("nan")), + (float("nan"), 100.0), + (1e6, 1e6 + 1), # True - within rtol + (1e6, 1e6 + 100), # False - outside rtol + (1e-6, 1.1e-6), # False - outside atol + (1e-7, 1.1e-7), # True - outside atol + (1e6 + 0j, 1e6 + 1j), # True - within rtol + (1e6 + 0j, 1e6 + 100j), # False - outside rtol + ], + ) + def test_basic(self, a: float, b: float, swap: bool, xp: ArrayNamespace): + if swap: + b, a = a, b + a_xp = xp.asarray(a) + b_xp = xp.asarray(b) + + assert_equal(isclose(a_xp, b_xp), xp.asarray(np.isclose(a, b))) + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + ar_np = a * np.arange(10) + br_np = b * np.arange(10) + ar_xp = xp.asarray(ar_np) + br_xp = xp.asarray(br_np) + + assert_equal(isclose(ar_xp, br_xp), xp.asarray(np.isclose(ar_np, br_np))) + + @pytest.mark.parametrize("dtype", ["float32", "int32"]) + def test_broadcast(self, dtype: str, xp: ArrayNamespace): + dtype = getattr(xp, dtype) + a = xp.asarray([1, 2, 3], dtype=dtype) + b = xp.asarray([[1], [5]], dtype=dtype) + actual = isclose(a, b) + expect = xp.asarray( + [[True, False, False], [False, False, False]], dtype=xp.bool + ) + + assert_equal(actual, expect) + + def test_some_inf(self, xp: ArrayNamespace): + a = xp.asarray([0.0, 1.0, xp.inf, xp.inf, xp.inf]) + b = xp.asarray([1e-9, 1.0, xp.inf, -xp.inf, 2.0]) + actual = isclose(a, b) + assert_equal(actual, xp.asarray([True, True, True, False, False])) + + def test_equal_nan(self, xp: ArrayNamespace): + a = xp.asarray([xp.nan, xp.nan, 1.0]) + b = xp.asarray([xp.nan, 1.0, xp.nan]) + assert_equal(isclose(a, b), xp.asarray([False, False, False])) + assert_equal(isclose(a, b, equal_nan=True), xp.asarray([True, False, False])) + + @pytest.mark.parametrize("dtype", ["float32", "complex64", "int32"]) + def test_tolerance(self, dtype: str, xp: ArrayNamespace): + dtype = getattr(xp, dtype) + a = xp.asarray([100, 100], dtype=dtype) + b = xp.asarray([101, 102], dtype=dtype) + assert_equal(isclose(a, b), xp.asarray([False, False])) + assert_equal(isclose(a, b, atol=1), xp.asarray([True, False])) + assert_equal(isclose(a, b, rtol=0.01), xp.asarray([True, False])) + + # Attempt to trigger division by 0 in rtol on int dtype + assert_equal(isclose(a, b, rtol=0), xp.asarray([False, False])) + assert_equal(isclose(a, b, atol=1, rtol=0), xp.asarray([True, False])) + + @pytest.mark.parametrize("dtype", ["int8", "uint8"]) + def test_tolerance_integer_overflow(self, dtype: str, xp: ArrayNamespace): + """1/rtol is too large for dtype""" + a = xp.asarray([100, 100], dtype=getattr(xp, dtype)) + b = xp.asarray([100, 101], dtype=getattr(xp, dtype)) + assert_equal(isclose(a, b), xp.asarray([True, False])) + + def test_very_small_numbers(self, xp: ArrayNamespace): + a = xp.asarray([1e-9, 1e-9]) + b = xp.asarray([1.0001e-9, 1.00001e-9]) + # Difference is below default atol + assert_equal(isclose(a, b), xp.asarray([True, True])) + # Use only rtol + assert_equal(isclose(a, b, atol=0), xp.asarray([False, True])) + assert_equal(isclose(a, b, atol=0, rtol=0), xp.asarray([False, False])) + + def test_bool_dtype(self, xp: ArrayNamespace): + a = xp.asarray([False, True, False]) + b = xp.asarray([True, True, False]) + assert_equal(isclose(a, b), xp.asarray([False, True, True])) + assert_equal(isclose(a, b, atol=1), xp.asarray([True, True, True])) + assert_equal(isclose(a, b, atol=2), xp.asarray([True, True, True])) + assert_equal(isclose(a, b, rtol=1), xp.asarray([True, True, True])) + assert_equal(isclose(a, b, rtol=2), xp.asarray([True, True, True])) + + # Test broadcasting + assert_equal( + isclose(a, xp.asarray(True), atol=1), xp.asarray([True, True, True]) + ) + assert_equal( + isclose(xp.asarray(True), b, atol=1), xp.asarray([True, True, True]) + ) + + @pytest.mark.skip_xp_backend(Backend.SPARSE, reason="index by sparse array") + @pytest.mark.skip_xp_backend(Backend.ARRAY_API_STRICTEST, reason="unknown shape") + def test_none_shape(self, xp: ArrayNamespace): + a = xp.asarray([1, 5, 0]) + b = xp.asarray([1, 4, 2]) + b = b[a < 5] + a = a[a < 5] + assert_equal(isclose(a, b), xp.asarray([True, False])) + + @pytest.mark.skip_xp_backend(Backend.SPARSE, reason="index by sparse array") + @pytest.mark.skip_xp_backend(Backend.ARRAY_API_STRICTEST, reason="unknown shape") + def test_none_shape_bool(self, xp: ArrayNamespace): + a = xp.asarray([True, True, False]) + b = xp.asarray([True, False, True]) + b = b[a] + a = a[a] + assert_equal(isclose(a, b), xp.asarray([True, False])) + + @pytest.mark.skip_xp_backend(Backend.NUMPY_READONLY, reason="xp=xp") + def test_python_scalar(self, xp: ArrayNamespace): + a = xp.asarray([0.0, 0.1], dtype=xp.float32) + assert_equal(isclose(a, 0.0), xp.asarray([True, False])) + assert_equal(isclose(0.0, a), xp.asarray([True, False])) + + a = xp.asarray([0, 1], dtype=xp.int16) + assert_equal(isclose(a, 0), xp.asarray([True, False])) + assert_equal(isclose(0, a), xp.asarray([True, False])) + + assert_equal(isclose(0, 0, xp=xp), xp.asarray(True)) + assert_equal(isclose(0, 1, xp=xp), xp.asarray(False)) + + def test_all_python_scalars(self): + with pytest.raises(TypeError, match=r"array_namespace requires .* array input"): + _ = isclose(0, 0) + + def test_xp(self, xp: ArrayNamespace): + a = xp.asarray([0.0, 0.0]) + b = xp.asarray([1e-9, 1e-4]) + assert_equal(isclose(a, b, xp=xp), xp.asarray([True, False])) + + @pytest.mark.parametrize("equal_nan", [True, False]) + def test_device(self, xp: ArrayNamespace, device: Device, equal_nan: bool): + a = xp.asarray([0.0, 0.0, xp.nan], device=device) + b = xp.asarray([1e-9, 1e-4, xp.nan], device=device) + res = isclose(a, b, equal_nan=equal_nan) + assert get_device(res) == device + + def test_array_on_device_with_scalar(self, xp: ArrayNamespace, device: Device): + a = xp.asarray([0.01, 0.5, 0.8, 0.9, 1.00001], device=device, dtype=xp.float64) + b = 1 + res = isclose(a, b) + assert get_device(res) == device + assert_equal(res, xp.asarray([False, False, False, False, True], device=device)) + + a = 0.1 + b = xp.asarray([0.01, 0.5, 0.8, 0.9, 0.100001], device=device, dtype=xp.float64) + res = isclose(a, b) + assert get_device(res) == device + assert_equal(res, xp.asarray([False, False, False, False, True], device=device)) + + +class TestNanToNum: + def test_bool(self, xp: ArrayNamespace) -> None: + a = xp.asarray([True]) + assert_equal(nan_to_num(a, xp=xp), a) + + def test_scalar_pos_inf(self, xp: ArrayNamespace, infinity: float) -> None: + a = xp.inf + assert_equal(nan_to_num(a, xp=xp), xp.asarray(infinity)) + + def test_scalar_neg_inf(self, xp: ArrayNamespace, infinity: float) -> None: + a = -xp.inf + assert_equal(nan_to_num(a, xp=xp), -xp.asarray(infinity)) + + def test_scalar_nan(self, xp: ArrayNamespace) -> None: + a = xp.nan + assert_equal(nan_to_num(a, xp=xp), xp.asarray(0.0)) + + def test_real(self, xp: ArrayNamespace, infinity: float) -> None: + a = xp.asarray([xp.inf, -xp.inf, xp.nan, -128, 128]) + assert_equal( + nan_to_num(a, xp=xp), + xp.asarray( + [ + infinity, + -infinity, + 0.0, + -128, + 128, + ] + ), + ) + + def test_complex(self, xp: ArrayNamespace, infinity: float) -> None: + a = xp.asarray( + [ + complex(xp.inf, xp.nan), + xp.nan, + complex(xp.nan, xp.inf), + ] + ) + assert_equal( + nan_to_num(a), + xp.asarray([complex(infinity, 0), complex(0, 0), complex(0, infinity)]), + ) + + def test_empty_array(self, xp: ArrayNamespace) -> None: + a = xp.asarray([], dtype=xp.float32) # forced dtype due to torch + assert_equal(nan_to_num(a, xp=xp), a) + assert xp.isdtype(nan_to_num(a, xp=xp).dtype, xp.float32) + + @pytest.mark.parametrize( + ("in_vals", "fill_value", "out_vals"), + [ + ([1, 2, np.nan, 4], 3, [1.0, 2.0, 3.0, 4.0]), + ([1, 2, np.nan, 4], 3.0, [1.0, 2.0, 3.0, 4.0]), + ( + [ + complex(1, 1), + complex(2, 2), + complex(np.nan, 0), + complex(4, 4), + ], + 3, + [ + complex(1.0, 1.0), + complex(2.0, 2.0), + complex(3.0, 0.0), + complex(4.0, 4.0), + ], + ), + ( + [ + complex(1, 1), + complex(2, 2), + complex(0, np.nan), + complex(4, 4), + ], + 3.0, + [ + complex(1.0, 1.0), + complex(2.0, 2.0), + complex(0.0, 3.0), + complex(4.0, 4.0), + ], + ), + ( + [ + complex(1, 1), + complex(2, 2), + complex(np.nan, np.nan), + complex(4, 4), + ], + 3.0, + [ + complex(1.0, 1.0), + complex(2.0, 2.0), + complex(3.0, 3.0), + complex(4.0, 4.0), + ], + ), + ], + ) + def test_fill_value_success( + self, + xp: ArrayNamespace, + in_vals: Array, + fill_value: float, + out_vals: Array, + ) -> None: + a = xp.asarray(in_vals) + assert_equal( + nan_to_num(a, fill_value=fill_value, xp=xp), + xp.asarray(out_vals), + ) + + def test_fill_value_failure(self, xp: ArrayNamespace) -> None: + a = xp.asarray( + [ + complex(1, 1), + complex(xp.nan, xp.nan), + complex(3, 3), + ] + ) + with pytest.raises( + TypeError, + match="Complex fill values are not supported", + ): + _ = nan_to_num( + a, + fill_value=complex(2, 2), # type: ignore[arg-type] # pyright: ignore[reportArgumentType] + xp=xp, + ) + + +class TestSinc: + def test_simple(self, xp: ArrayNamespace): + assert_equal(sinc(xp.asarray(0.0)), xp.asarray(1.0)) + x = xp.asarray(np.linspace(-1, 1, 100)) + w = sinc(x) + # check symmetry + assert_close(w, xp.flip(w, axis=0)) + + @pytest.mark.parametrize("x", [0, 1 + 3j]) + def test_dtype(self, xp: ArrayNamespace, x: complex): + with pytest.raises(ValueError, match="real floating data type"): + _ = sinc(xp.asarray(x)) + + def test_3d(self, xp: ArrayNamespace): + x = np.arange(18, dtype=np.float64).reshape((3, 3, 2)) + expected = np.zeros_like(x) + expected[0, 0, 0] = 1 + x = xp.asarray(x) + expected = xp.asarray(expected) + assert_close(sinc(x), expected, atol=1e-15) + + def test_device(self, xp: ArrayNamespace, device: Device): + x = xp.asarray(0.0, device=device) + assert get_device(sinc(x)) == device + + def test_xp(self, xp: ArrayNamespace): + assert_equal(sinc(xp.asarray(0.0), xp=xp), xp.asarray(1.0)) + + +class TestAngle: + def test_simple(self, xp: ArrayNamespace): + a = xp.asarray([1, 0]) + res = angle(a) + expected = xp.asarray([0.0, 0.0], dtype=res.dtype) + assert_equal(res, expected) + + def test_basic(self, xp: ArrayNamespace): + x = xp.asarray( + [ + 1 + 3j, + np.sqrt(2) / 2.0 + 1j * np.sqrt(2) / 2, + 1, + 1j, + -1, + -1j, + 1 - 3j, + -1 + 3j, + ], + dtype=xp.complex128, + ) + expected = xp.asarray( + [ + np.arctan(3.0 / 1.0), + np.arctan(1.0), + 0, + np.pi / 2, + np.pi, + -np.pi / 2.0, + -np.arctan(3.0 / 1.0), + np.pi - np.arctan(3.0 / 1.0), + ], + dtype=xp.float64, + ) + assert_close(angle(x, xp=xp), expected, rtol=0, atol=1e-11) + assert_close( + angle(x, deg=True, xp=xp), + expected * 180 / xp.pi, + rtol=0, + atol=1e-11, + ) + + def test_real(self, xp: ArrayNamespace): + x = xp.asarray([0.0, -0.0, 1.0, -1.0]) + expected = xp.asarray([0.0, xp.pi, 0.0, xp.pi], dtype=x.dtype) + assert_close(angle(x, xp=xp), expected) + + def test_complex(self, xp: ArrayNamespace): + a = xp.asarray([1 + 1j, 1 - 1j, -1 + 1j, -1 - 1j]) + expected = xp.asarray([xp.pi / 4, -xp.pi / 4, 3 * xp.pi / 4, -3 * xp.pi / 4]) + res = angle(a, xp=xp) + assert_equal(res, expected) + + def test_integral(self, xp: ArrayNamespace): + x = xp.asarray([0, -1, 1], dtype=xp.int32) + actual = angle(x, xp=xp) + expected = xp.asarray( + [0.0, xp.pi, 0.0], dtype=default_dtype(xp, device=get_device(x)) + ) + assert_close(actual, expected) + + def test_2d(self, xp: ArrayNamespace): + a = xp.asarray([[1 + 1j, 1 - 1j], [-1 + 1j, -1 - 1j]]) + expected = xp.asarray( + [[xp.pi / 4, -xp.pi / 4], [3 * xp.pi / 4, -3 * xp.pi / 4]] + ) + res = angle(a, xp=xp) + assert_equal(res, expected) + + @pytest.mark.skip_xp_backend(Backend.TORCH, reason="materialize 'meta' device") + def test_device(self, xp: ArrayNamespace, device: Device): + a = xp.asarray([1 + 1j], device=device) + assert get_device(angle(a)) == device + + +class TestDeg2Rad: + def test_basic(self, xp: ArrayNamespace): + x = xp.asarray([0.0, 90.0, 180.0, 270.0, 360.0]) + expected = xp.asarray([0.0, xp.pi / 2, xp.pi, 3 * xp.pi / 2, 2 * xp.pi]) + assert_close(deg2rad(x), expected) + + @pytest.mark.parametrize("dtype_name", ["int32", "int64"]) + def test_integral(self, xp: ArrayNamespace, dtype_name: str): + x = xp.asarray([0, 90, 180], dtype=getattr(xp, dtype_name)) + actual = deg2rad(x, xp=xp) + expected = xp.asarray( + [0.0, xp.pi / 2, xp.pi], dtype=default_dtype(xp, device=get_device(x)) + ) + assert actual.dtype == expected.dtype + assert_close(actual, expected) + + def test_complex(self, xp: ArrayNamespace): + x = xp.asarray([180 + 90j], dtype=xp.complex64) + actual = deg2rad(x, xp=xp) + expected = xp.asarray([xp.pi + xp.pi / 2 * 1j], dtype=x.dtype) + assert actual.dtype == x.dtype + assert_close(actual, expected) + + def test_bool(self, xp: ArrayNamespace): + x = xp.asarray([True]) + with pytest.raises(TypeError, match="integral, real floating, or complex"): + _ = deg2rad(x, xp=xp) + + def test_device(self, xp: ArrayNamespace, device: Device): + x = xp.asarray([0.0, 90.0, 180.0], device=device) + assert get_device(deg2rad(x)) == device + + +class TestRad2Deg: + def test_basic(self, xp: ArrayNamespace): + x = xp.asarray([0.0, xp.pi / 2, xp.pi, 3 * xp.pi / 2, 2 * xp.pi]) + expected = xp.asarray([0.0, 90.0, 180.0, 270.0, 360.0]) + assert_close(rad2deg(x), expected) + + @pytest.mark.parametrize("dtype_name", ["int32", "int64"]) + def test_integral(self, xp: ArrayNamespace, dtype_name: str): + x = xp.asarray([0, 1, 2], dtype=getattr(xp, dtype_name)) + actual = rad2deg(x, xp=xp) + expected = xp.asarray( + [0.0, 180 / xp.pi, 360 / xp.pi], + dtype=default_dtype(xp, device=get_device(x)), + ) + assert actual.dtype == expected.dtype + assert_close(actual, expected) + + def test_complex(self, xp: ArrayNamespace): + x = xp.asarray([xp.pi + xp.pi / 2 * 1j], dtype=xp.complex64) + actual = rad2deg(x, xp=xp) + expected = xp.asarray([180 + 90j], dtype=x.dtype) + assert actual.dtype == x.dtype + assert_close(actual, expected) + + def test_bool(self, xp: ArrayNamespace): + x = xp.asarray([True]) + with pytest.raises(TypeError, match="integral, real floating, or complex"): + _ = rad2deg(x, xp=xp) + + def test_device(self, xp: ArrayNamespace, device: Device): + x = xp.asarray([0.0, xp.pi / 2, xp.pi], device=device) + assert get_device(rad2deg(x)) == device diff --git a/tests/test_funcs.py b/tests/test_funcs.py deleted file mode 100644 index 090d592a..00000000 --- a/tests/test_funcs.py +++ /dev/null @@ -1,2736 +0,0 @@ -import inspect -import math -import warnings -from collections.abc import Callable -from typing import Any, Literal, cast - -import hypothesis -import hypothesis.extra.numpy as npst -import numpy as np -import pytest -from hypothesis import given -from hypothesis import strategies as st -from typing_extensions import override - -import array_api_extra._delegation as delegated_func -from array_api_extra import ( - angle, - apply_where, - argpartition, - at, - atleast_nd, - broadcast_shapes, - cov, - create_diagonal, - default_dtype, - deg2rad, - diag_indices, - expand_dims, - isclose, - isin, - kron, - nan_to_num, - nanmax, - nanmin, - nansum, - nunique, - one_hot, - pad, - partition, - rad2deg, - setdiff1d, - sinc, - tril_indices, - triu_indices, - union1d, - unravel_index, -) -from array_api_extra import ( - searchsorted as xpx_searchsorted, -) -from array_api_extra._lib import _funcs as functions -from array_api_extra._lib._backends import NUMPY_VERSION, Backend -from array_api_extra._lib._funcs import searchsorted as _funcs_searchsorted -from array_api_extra._lib._utils._compat import ( - array_namespace, - is_jax_namespace, - is_torch_namespace, -) -from array_api_extra._lib._utils._compat import device as get_device -from array_api_extra._lib._utils._helpers import eager_shape, ndindex -from array_api_extra._lib._utils._typing import Array, ArrayNamespace, Device -from array_api_extra.testing import assert_close, assert_equal, lazy_xp_function - -lazy_xp_function(apply_where) -lazy_xp_function(argpartition) -lazy_xp_function(atleast_nd) -lazy_xp_function(broadcast_shapes) -lazy_xp_function(cov) -lazy_xp_function(create_diagonal) -lazy_xp_function(default_dtype) -lazy_xp_function(deg2rad) -lazy_xp_function(diag_indices) -lazy_xp_function(expand_dims) -lazy_xp_function(isclose) -lazy_xp_function(isin) -lazy_xp_function(kron) -lazy_xp_function(nan_to_num) -lazy_xp_function(nansum) -lazy_xp_function(nunique) -lazy_xp_function(one_hot) -lazy_xp_function(pad) -lazy_xp_function(partition) -lazy_xp_function(rad2deg) -# FIXME calls in1d which calls xp.unique_values without size -lazy_xp_function(setdiff1d, jax_jit=False) -lazy_xp_function(sinc) -lazy_xp_function(tril_indices) -lazy_xp_function(triu_indices) -lazy_xp_function(union1d, jax_jit=False) -lazy_xp_function(xpx_searchsorted) -lazy_xp_function(_funcs_searchsorted) - - -def test_all_contains_all_public_functions(): - public_functions = { - name - for name, obj in inspect.getmembers(functions, inspect.isfunction) - if not name.startswith("_") and obj.__module__ == functions.__name__ - } - missing = sorted(public_functions - set(functions.__all__)) - extra = sorted(set(functions.__all__) - public_functions) - assert public_functions == set(functions.__all__), ( - f"Missing from __all__: {missing}\tExtra in __all__: {extra}" - ) - public_functions = { - name - for name, obj in inspect.getmembers(delegated_func, inspect.isfunction) - if not name.startswith("_") and obj.__module__ == delegated_func.__name__ - } - missing = sorted(public_functions - set(delegated_func.__all__)) - extra = sorted(set(delegated_func.__all__) - public_functions) - assert public_functions == set(delegated_func.__all__), ( - f"Missing from __all__: {missing}\tExtra in __all__: {extra}" - ) - - -class TestApplyWhere: - @staticmethod - def f1(x: Array, y: Array | int = 10) -> Array: - return x + y - - @staticmethod - def f2(x: Array, y: Array | int = 10) -> Array: - return x - y - - def test_f1_f2(self, xp: ArrayNamespace): - x = xp.asarray([1, 2, 3, 4]) - cond = x % 2 == 0 - actual = apply_where(cond, x, self.f1, self.f2) - expect = xp.where(cond, self.f1(x), self.f2(x)) - assert_equal(actual, expect) - - def test_fill_value(self, xp: ArrayNamespace): - x = xp.asarray([1, 2, 3, 4]) - cond = x % 2 == 0 - actual = apply_where(x % 2 == 0, x, self.f1, fill_value=0) - expect = xp.where(cond, self.f1(x), xp.asarray(0)) - assert_equal(actual, expect) - - actual = apply_where(x % 2 == 0, x, self.f1, fill_value=xp.asarray(0)) - assert_equal(actual, expect) - - def test_args_tuple(self, xp: ArrayNamespace): - x = xp.asarray([1, 2, 3, 4]) - y = xp.asarray([10, 20, 30, 40]) - cond = x % 2 == 0 - actual = apply_where(cond, (x, y), self.f1, self.f2) - expect = xp.where(cond, self.f1(x, y), self.f2(x, y)) - assert_equal(actual, expect) - - def test_broadcast(self, xp: ArrayNamespace): - x = xp.asarray([1, 2]) - y = xp.asarray([[10], [20], [30]]) - cond = xp.broadcast_to(xp.asarray(True), (4, 1, 1)) - - actual = apply_where(cond, (x, y), self.f1, self.f2) - expect = xp.where(cond, self.f1(x, y), self.f2(x, y)) - assert_equal(actual, expect) - - actual = apply_where( - cond, - (x, y), - lambda x, _: x, - lambda _, y: y, - ) - expect = xp.where(cond, x, y) - assert_equal(actual, expect) - - # Shaped fill_value - actual = apply_where(cond, x, self.f1, fill_value=y) - expect = xp.where(cond, self.f1(x), y) - assert_equal(actual, expect) - - def test_dtype_propagation(self, xp: ArrayNamespace, library: Backend): - x = xp.asarray([1, 2], dtype=xp.int8) - y = xp.asarray([3, 4], dtype=xp.int16) - cond = x % 2 == 0 - - mxp = np if library is Backend.DASK else xp - actual = apply_where( - cond, - (x, y), - self.f1, - lambda x, y: mxp.astype(x - y, xp.int64), # pyright: ignore[reportArgumentType] # pyrefly: ignore[bad-argument-type] - ) - assert actual.dtype == xp.int64 - - actual = apply_where(cond, y, self.f1, fill_value=5) - assert actual.dtype == xp.int16 - - @pytest.mark.parametrize("fill_value_raw", [3, [3, 4]]) - @pytest.mark.parametrize( - ("fill_value_dtype", "expect_dtype"), [("int32", "int32"), ("int8", "int16")] - ) - def test_dtype_propagation_fill_value( - self, - xp: ArrayNamespace, - fill_value_raw: int | list[int], - fill_value_dtype: str, - expect_dtype: str, - ): - x = xp.asarray([1, 2], dtype=xp.int16) - cond = x % 2 == 0 - fill_value = xp.asarray(fill_value_raw, dtype=getattr(xp, fill_value_dtype)) - - actual = apply_where(cond, x, self.f1, fill_value=fill_value) - assert actual.dtype == getattr(xp, expect_dtype) - - def test_dont_overwrite_fill_value(self, xp: ArrayNamespace): - x = xp.asarray([1, 2]) - fill_value = xp.asarray([100, 200]) - actual = apply_where(x % 2 == 0, x, self.f1, fill_value=fill_value) - assert_equal(actual, xp.asarray([100, 12])) - assert_equal(fill_value, xp.asarray([100, 200])) - - @pytest.mark.skip_xp_backend( - Backend.ARRAY_API_STRICTEST, - reason="no boolean indexing -> run everywhere", - ) - @pytest.mark.skip_xp_backend( - Backend.SPARSE, - reason="no indexing by sparse array -> run everywhere", - ) - def test_dont_run_on_false(self, xp: ArrayNamespace): - x = xp.asarray([1.0, 2.0, 0.0]) - y = xp.asarray([0.0, 3.0, 4.0]) - # On NumPy, division by zero will trigger warnings - actual = apply_where( - x == 0, - (x, y), - lambda x, y: x / y, - lambda x, y: y / x, - ) - assert_equal(actual, xp.asarray([0.0, 1.5, 0.0])) - - def test_bad_args(self, xp: ArrayNamespace): - x = xp.asarray([1, 2, 3, 4]) - cond = x % 2 == 0 - # Neither f2 nor fill_value - with pytest.raises(TypeError, match="Exactly one of"): - apply_where(cond, x, self.f1) # type: ignore[call-overload] # pyright: ignore[reportCallIssue] - # Both f2 and fill_value - with pytest.raises(TypeError, match="Exactly one of"): - apply_where(cond, x, self.f1, self.f2, fill_value=0) # type: ignore[call-overload] # pyright: ignore[reportCallIssue] - - @pytest.mark.skip_xp_backend(Backend.NUMPY_READONLY, reason="xp=xp") - def test_xp(self, xp: ArrayNamespace): - x = xp.asarray([1, 2, 3, 4]) - cond = x % 2 == 0 - actual = apply_where(cond, x, self.f1, self.f2, xp=xp) - expect = xp.where(cond, self.f1(x), self.f2(x)) - assert_equal(actual, expect) - - def test_device(self, xp: ArrayNamespace, device: Device): - x = xp.asarray([1, 2, 3, 4], device=device) - y = apply_where(x % 2 == 0, x, self.f1, self.f2) - assert get_device(y) == device - y = apply_where(x % 2 == 0, x, self.f1, fill_value=0) - assert get_device(y) == device - y = apply_where(x % 2 == 0, x, self.f1, fill_value=x) - assert get_device(y) == device - - @pytest.mark.filterwarnings("ignore::RuntimeWarning") # overflows, etc. - @hypothesis.settings( - # The xp and library fixtures are not regenerated between hypothesis iterations - suppress_health_check=[hypothesis.HealthCheck.function_scoped_fixture], - # JAX can take a long time to initialize on the first call - deadline=None, - ) - @given( - n_arrays=st.integers(min_value=0, max_value=3), - n_kwarrays=st.integers(min_value=0, max_value=3), - rng_seed=st.integers(min_value=1000000000, max_value=9999999999), - dtype=npst.floating_dtypes(sizes=(32, 64)), - p=st.floats(min_value=0, max_value=1), - data=st.data(), - ) - def test_hypothesis( - self, - n_arrays: int, - n_kwarrays: int, - rng_seed: int, - dtype: np.dtype[Any], - p: float, - data: st.DataObject, - xp: ArrayNamespace, - library: Backend, - ): - if ( - library.like(Backend.NUMPY) - and NUMPY_VERSION < (2, 0) - and dtype.type is np.float32 - ): - pytest.xfail(reason="NumPy 1.x dtype promotion for scalars") - - _ = hypothesis.assume(n_arrays + n_kwarrays > 0) - mbs = npst.mutually_broadcastable_shapes( - num_shapes=1 + n_arrays + n_kwarrays, min_side=0 - ) - input_shapes, _ = data.draw(mbs) - cond_shape = input_shapes[0] - shapes = input_shapes[1 : 1 + n_arrays] - kwshapes = input_shapes[1 + n_arrays :] - - # cupy/cupy#8382 - # https://github.com/jax-ml/jax/issues/26658 - elements = {"allow_subnormal": not library.like(Backend.CUPY, Backend.JAX)} - - fill_value = xp.asarray( - data.draw(npst.arrays(dtype=dtype.type, shape=(), elements=elements)) - ) - float_fill_value = float(fill_value) - if library is Backend.CUPY and dtype.type is np.float32: - # Avoid data-dependent dtype promotion when encountering subnormals - # close to the max float32 value - float_fill_value = float(np.clip(float_fill_value, -1e38, 1e38)) - - arrays = tuple( - xp.asarray( - data.draw(npst.arrays(dtype=dtype.type, shape=shape, elements=elements)) - ) - for shape in shapes - ) - - kwargs = { - f"kw{n}": xp.asarray( - data.draw(npst.arrays(dtype=dtype.type, shape=shape, elements=elements)) - ) - for n, shape in enumerate(kwshapes) - } - kwkeys = kwargs.keys() - - def f1(*args: Array, **kwargs: dict[str, Array]) -> Array: - assert kwargs.keys() == kwkeys - args_kwargs = cast(tuple[Array, ...], (*args, *kwargs.values())) - return cast(Array, sum(args_kwargs)) - - def f2(*args: Array, **kwargs: dict[str, Array]) -> Array: - assert kwargs.keys() == kwkeys - args_kwargs = cast(tuple[Array, ...], (*args, *kwargs.values())) - return cast(Array, sum(args_kwargs) / 2) - - rng = np.random.default_rng(rng_seed) - cond = xp.asarray(rng.random(size=cond_shape) > p) - - res1 = apply_where(cond, arrays, f1, fill_value=fill_value, kwargs=kwargs) - res2 = apply_where(cond, arrays, f1, f2, kwargs=kwargs) - res3 = apply_where(cond, arrays, f1, fill_value=float_fill_value, kwargs=kwargs) - - ref1 = xp.where(cond, f1(*arrays, **kwargs), fill_value) - ref2 = xp.where(cond, f1(*arrays, **kwargs), f2(*arrays, **kwargs)) - ref3 = xp.where(cond, f1(*arrays, **kwargs), float_fill_value) - - assert_close(res1, ref1, rtol=2e-16) - assert_equal(res2, ref2) - assert_equal(res3, ref3) - - -class TestAtLeastND: - def test_0D(self, xp: ArrayNamespace): - x = xp.asarray(1.0) - - y = atleast_nd(x, ndim=0) - assert_equal(y, x) - - y = atleast_nd(x, ndim=1) - assert_equal(y, xp.ones((1,))) - - y = atleast_nd(x, ndim=5) - assert_equal(y, xp.ones((1, 1, 1, 1, 1))) - - @pytest.mark.parametrize( - ("input_shape", "ndim", "expected_shape"), - [ - ((1,), 0, (1,)), - ((5,), 1, (5,)), - ((2,), 2, (1, 2)), - ((3,), 3, (1, 1, 3)), - ((2,), 5, (1, 1, 1, 1, 2)), - ], - ) - def test_1D_shapes( - self, - input_shape: tuple[int], - ndim: int, - expected_shape: tuple[int], - xp: ArrayNamespace, - ): - n = math.prod(input_shape) - x = xp.asarray(np.arange(n).reshape(input_shape)) - y = atleast_nd(x, ndim=ndim) - - assert y.shape == expected_shape - assert xp.sum(y) == int(n * (n - 1) / 2) - - def test_1D_values(self, xp: ArrayNamespace): - x = xp.asarray([0, 1]) - - y = atleast_nd(x, ndim=0) - assert_equal(y, x) - - y = atleast_nd(x, ndim=1) - assert_equal(y, x) - - y = atleast_nd(x, ndim=2) - assert_equal(y, xp.asarray([[0, 1]])) - - y = atleast_nd(x, ndim=5) - assert_equal(y, xp.asarray([[[[[0, 1]]]]])) - - @pytest.mark.parametrize( - ("input_shape", "ndim", "expected_shape"), - [ - ((2, 1), 0, (2, 1)), - ((5, 2), 1, (5, 2)), - ((2, 1), 2, (2, 1)), - ((3, 1), 3, (1, 3, 1)), - ((2, 8), 5, (1, 1, 1, 2, 8)), - ], - ) - def test_2D_shapes( - self, - input_shape: tuple[int], - ndim: int, - expected_shape: tuple[int], - xp: ArrayNamespace, - ): - n = math.prod(input_shape) - x = xp.asarray(np.arange(n).reshape(input_shape)) - y = atleast_nd(x, ndim=ndim) - - assert y.shape == expected_shape - assert xp.sum(y) == int(n * (n - 1) / 2) - - def test_2D_values(self, xp: ArrayNamespace): - x = xp.asarray([[3.0], [4.0]]) - - y = atleast_nd(x, ndim=0) - assert_equal(y, x) - - y = atleast_nd(x, ndim=2) - assert_equal(y, x) - - y = atleast_nd(x, ndim=3) - assert_equal(y, xp.asarray([[[3.0], [4.0]]])) - - y = atleast_nd(x, ndim=5) - assert_equal(y, xp.asarray([[[[[3.0], [4.0]]]]])) - - @pytest.mark.parametrize( - ("input_shape", "ndim", "expected_shape"), - [ - ((2, 1, 1), 0, (2, 1, 1)), - ((1, 5, 2), 1, (1, 5, 2)), - ((2, 1, 1), 2, (2, 1, 1)), - ((1, 3, 1), 3, (1, 3, 1)), - ((2, 8, 1), 5, (1, 1, 2, 8, 1)), - ], - ) - def test_3D_shapes( - self, - input_shape: tuple[int], - ndim: int, - expected_shape: tuple[int], - xp: ArrayNamespace, - ): - n = math.prod(input_shape) - x = xp.asarray(np.arange(n).reshape(input_shape)) - y = atleast_nd(x, ndim=ndim) - - assert y.shape == expected_shape - assert xp.sum(y) == int(n * (n - 1) / 2) - - def test_3D_values(self, xp: ArrayNamespace): - x = xp.asarray([[[3.0], [2.0]]]) - - y = atleast_nd(x, ndim=0) - assert_equal(y, x) - - y = atleast_nd(x, ndim=2) - assert_equal(y, x) - - y = atleast_nd(x, ndim=3) - assert_equal(y, x) - - y = atleast_nd(x, ndim=5) - assert_equal(y, xp.asarray([[[[[3.0], [2.0]]]]])) - - @pytest.mark.parametrize( - ("input_shape", "ndim", "expected_shape"), - [ - ((2, 1, 1, 2, 1), 0, (2, 1, 1, 2, 1)), - ((1, 5, 2, 3, 2), 2, (1, 5, 2, 3, 2)), - ((2, 1, 1, 5, 2), 5, (2, 1, 1, 5, 2)), - ((1, 3, 1, 2, 1), 6, (1, 1, 3, 1, 2, 1)), - ((2, 8, 1, 9, 8), 9, (1, 1, 1, 1, 2, 8, 1, 9, 8)), - ], - ) - def test_5D_shapes( - self, - input_shape: tuple[int], - ndim: int, - expected_shape: tuple[int], - xp: ArrayNamespace, - ): - n = math.prod(input_shape) - x = xp.asarray(np.arange(n).reshape(input_shape)) - y = atleast_nd(x, ndim=ndim) - - assert y.shape == expected_shape - assert xp.sum(y) == int(n * (n - 1) / 2) - - def test_5D_values(self, xp: ArrayNamespace): - x = xp.asarray([[[[[3.0]], [[2.0]]]]]) - - y = atleast_nd(x, ndim=0) - assert_equal(y, x) - - y = atleast_nd(x, ndim=4) - assert_equal(y, x) - - y = atleast_nd(x, ndim=5) - assert_equal(y, x) - - y = atleast_nd(x, ndim=6) - assert_equal(y, xp.asarray([[[[[[3.0]], [[2.0]]]]]])) - - y = atleast_nd(x, ndim=9) - assert_equal(y, xp.asarray([[[[[[[[[3.0]], [[2.0]]]]]]]]])) - - -@pytest.mark.filterwarnings("ignore:.*removed in v1.0.0.*:DeprecationWarning") -class TestBroadcastShapes: - def test_delegates_known_integer_shapes(self, monkeypatch: pytest.MonkeyPatch): - calls = [] - - def mock_broadcast_shapes(*shapes: tuple[int, ...]) -> tuple[int, ...]: - calls.append(shapes) - return (99,) - - monkeypatch.setattr(np, "broadcast_shapes", mock_broadcast_shapes) - - assert broadcast_shapes((2,), (1,), xp=np) == (99,) - assert calls == [((2,), (1,))] - - def test_fallback_without_xp(self, monkeypatch: pytest.MonkeyPatch): - def mock_broadcast_shapes(*_shapes: tuple[int, ...]) -> tuple[int, ...]: - msg = "Native delegation should not be used without xp" - raise AssertionError(msg) - - monkeypatch.setattr(np, "broadcast_shapes", mock_broadcast_shapes) - - assert broadcast_shapes((2,), (1,)) == (2,) - - @pytest.mark.skip_xp_backend(Backend.NUMPY_READONLY, reason="xp=xp") - def test_xp(self, xp: ArrayNamespace): - assert broadcast_shapes((2, 3), (2, 1), xp=xp) == (2, 3) - - @pytest.mark.parametrize( - "args", - [ - (), - ((),), - ((), ()), - ((1,),), - ((1,), (1,)), - ((2,), (1,)), - ((3, 1, 4), (2, 1)), - ((1, 1, 4), (2, 1)), - ((1,), ()), - ((), (2,), ()), - ((0,),), - ((0,), (1,)), - ((2, 0), (1, 1)), - ((2, 0, 3), (2, 1, 1)), - ], - ) - def test_simple(self, args: tuple[tuple[int, ...], ...]): - expect = np.broadcast_shapes(*args) - actual = broadcast_shapes(*args) - assert actual == expect - - @pytest.mark.parametrize( - "args", - [ - ((2,), (3,)), - ((2, 3), (1, 2)), - ((2,), (0,)), - ((2, 0, 2), (1, 3, 1)), - ], - ) - def test_fail(self, args: tuple[tuple[int, ...], ...]): - match = "cannot be broadcast to a single shape" - with pytest.raises(ValueError, match=match): - _ = np.broadcast_shapes(*args) - with pytest.raises(ValueError, match=match): - _ = broadcast_shapes(*args) - - @pytest.mark.parametrize( - "args", - [ - ((None,), (None,)), - ((math.nan,), (None,)), - ((1, None, 2, 4), (2, 3, None, 1), (2, None, None, 4)), - ((1, math.nan, 2), (4, 2, 3, math.nan), (4, 2, None, None)), - ((math.nan, 1), (None, 2), (None, 2)), - ], - ) - def test_none(self, args: tuple[tuple[float | None, ...], ...]): - expect = args[-1] - actual = broadcast_shapes(*args[:-1]) - assert actual == expect - - -class TestCov: - def test_basic(self, xp: ArrayNamespace): - assert_close( - cov(xp.asarray([[0, 2], [1, 1], [2, 0]], dtype=xp.float64).T), - xp.asarray([[1.0, -1.0], [-1.0, 1.0]], dtype=xp.float64), - ) - - def test_complex(self, xp: ArrayNamespace): - actual = cov(xp.asarray([[1, 2, 3], [1j, 2j, 3j]], dtype=xp.complex128)) - expect = xp.asarray([[1.0, -1.0j], [1.0j, 1.0]], dtype=xp.complex128) - assert_close(actual, expect) - - def test_complex_with_weights(self, xp: ArrayNamespace): - m = np.asarray( - [[1 + 1j, 2 + 2j, 4 + 1j], [3 - 1j, 5 + 2j, 7 + 0j]], - dtype=np.complex128, - ) - weights = np.asarray([1.0, 2.0, 1.0]) - correction = 0.5 # Force the generic implementation. - - weight_sum = weights.sum() - avg = (m * weights).sum(axis=-1, keepdims=True) / weight_sum - centered = m - avg - normalizer = weight_sum - correction * (weights**2).sum() / weight_sum - expected = (centered * weights) @ centered.conj().T / normalizer - - actual = cov( - xp.asarray(m), - correction=correction, - aweights=xp.asarray(weights), - ) - assert_close(actual, xp.asarray(expected)) - - def test_empty(self, xp: ArrayNamespace): - with warnings.catch_warnings(record=True): - warnings.simplefilter("always", RuntimeWarning) - warnings.simplefilter("always", UserWarning) - assert_equal( - cov(xp.asarray([], dtype=xp.float64)), - xp.asarray(xp.nan, dtype=xp.float64), - ) - assert_equal( - cov(xp.reshape(xp.asarray([], dtype=xp.float64), (0, 2))), - xp.reshape(xp.asarray([], dtype=xp.float64), (0, 0)), - ) - assert_equal( - cov(xp.reshape(xp.asarray([], dtype=xp.float64), (2, 0))), - xp.asarray([[xp.nan, xp.nan], [xp.nan, xp.nan]], dtype=xp.float64), - ) - - def test_combination(self, xp: ArrayNamespace): - x = xp.asarray([-2.1, -1, 4.3], dtype=xp.float64) - y = xp.asarray([3, 1.1, 0.12], dtype=xp.float64) - X = xp.stack((x, y), axis=0) - desired = xp.asarray([[11.71, -4.286], [-4.286, 2.144133]], dtype=xp.float64) - assert_close(cov(X), desired, rtol=1e-6) - assert_close(cov(x), xp.asarray(11.71, dtype=xp.float64)) - assert_close(cov(y), xp.asarray(2.144133, dtype=xp.float64), rtol=1e-6) - - @pytest.mark.xfail_xp_backend( - Backend.TORCH, reason="torch.cov does not support tensors on meta device" - ) - def test_device(self, xp: ArrayNamespace, device: Device): - x = xp.asarray([1, 2, 3], device=device) - assert get_device(cov(x)) == device - - @pytest.mark.skip_xp_backend(Backend.NUMPY_READONLY, reason="xp=xp") - def test_xp(self, xp: ArrayNamespace): - assert_close( - cov( - xp.asarray([[0.0, 2.0], [1.0, 1.0], [2.0, 0.0]], dtype=xp.float64).T, - xp=xp, - ), - xp.asarray([[1.0, -1.0], [-1.0, 1.0]], dtype=xp.float64), - ) - - def test_batch(self, xp: ArrayNamespace): - rng = np.random.default_rng(8847643423) - batch_shape = (3, 4) - n_var, n_obs = 3, 20 - m = rng.random((*batch_shape, n_var, n_obs)) - res = cov(xp.asarray(m)) - ref_list = [np.cov(m_) for m_ in np.reshape(m, (-1, n_var, n_obs))] - ref = np.reshape(np.stack(ref_list), (*batch_shape, n_var, n_var)) - assert_close(res, xp.asarray(ref)) - - @pytest.mark.parametrize("bias", [True, False, 0, 1]) - def test_bias(self, xp: ArrayNamespace, bias: bool): - # `bias` maps to `correction`: bias=True -> correction=0, bias=False -> 1. - x = np.array([-2.1, -1, 4.3]) - y = np.array([3, 1.1, 0.12]) - X = np.stack((x, y), axis=0) - ref = np.cov(X, bias=bias) - assert_close( - cov(xp.asarray(X, dtype=xp.float64), correction=0 if bias else 1), - xp.asarray(ref, dtype=xp.float64), - rtol=1e-6, - ) - - @pytest.mark.parametrize("bias", [True, False, 0, 1]) - def test_bias_batch(self, xp: ArrayNamespace, bias: bool): - rng = np.random.default_rng(8847643423) - batch_shape = (3, 4) - n_var, n_obs = 3, 20 - m = rng.random((*batch_shape, n_var, n_obs)) - res = cov(xp.asarray(m), correction=0 if bias else 1) - ref_list = [np.cov(m_, bias=bias) for m_ in np.reshape(m, (-1, n_var, n_obs))] - ref = np.reshape(np.stack(ref_list), (*batch_shape, n_var, n_var)) - assert_close(res, xp.asarray(ref)) - - def test_correction(self, xp: ArrayNamespace): - rng = np.random.default_rng(20260417) - m = rng.random((3, 20)) - for correction in (0, 1, 2): - ref = np.cov(m, ddof=correction) - res = cov(xp.asarray(m), correction=correction) - assert_close(res, xp.asarray(ref)) - - def test_correction_float(self, xp: ArrayNamespace): - # Float correction: reference computed by hand (numpy.cov rejects - # non-integer ddof; our generic path supports it). - rng = np.random.default_rng(20260417) - m = rng.random((3, 20)) - n = m.shape[-1] - centered = m - m.mean(axis=-1, keepdims=True) - ref = centered @ centered.T / (n - 1.5) - res = cov(xp.asarray(m), correction=1.5) - assert_close(res, xp.asarray(ref)) - - def test_axis(self, xp: ArrayNamespace): - rng = np.random.default_rng(20260417) - m = rng.random((20, 3)) # observations on axis 0 - ref = np.cov(m, rowvar=False) - res = cov(xp.asarray(m), axis=0) - assert_close(res, xp.asarray(ref)) - res_neg = cov(xp.asarray(m), axis=-2) - assert_close(res_neg, xp.asarray(ref)) - - def test_frequency_weights(self, xp: ArrayNamespace): - rng = np.random.default_rng(20260417) - m = rng.random((3, 10)) - fw = np.asarray([1, 2, 1, 3, 1, 2, 1, 1, 2, 1], dtype=np.int64) - ref = np.cov(m, fweights=fw) - res = cov(xp.asarray(m), fweights=xp.asarray(fw)) - assert_close(res, xp.asarray(ref)) - - def test_weights(self, xp: ArrayNamespace): - rng = np.random.default_rng(20260417) - m = rng.random((3, 10)) - aw = rng.random(10) - ref = np.cov(m, aweights=aw) - res = cov(xp.asarray(m), aweights=xp.asarray(aw)) - assert_close(res, xp.asarray(ref)) - - def test_both_weights(self, xp: ArrayNamespace): - rng = np.random.default_rng(20260417) - m = rng.random((3, 10)) - fw = np.asarray([1, 2, 1, 3, 1, 2, 1, 1, 2, 1], dtype=np.int64) - aw = rng.random(10) - for correction in (0, 1, 2): - ref = np.cov(m, ddof=correction, fweights=fw, aweights=aw) - res = cov( - xp.asarray(m), - correction=correction, - fweights=xp.asarray(fw), - aweights=xp.asarray(aw), - ) - assert_close(res, xp.asarray(ref)) - - def test_batch_with_weights(self, xp: ArrayNamespace): - rng = np.random.default_rng(20260417) - batch_shape = (2, 3) - n_var, n_obs = 3, 15 - m = rng.random((*batch_shape, n_var, n_obs)) - aw = rng.random(n_obs) - res = cov(xp.asarray(m), aweights=xp.asarray(aw)) - ref_list = [np.cov(m_, aweights=aw) for m_ in np.reshape(m, (-1, n_var, n_obs))] - ref = np.reshape(np.stack(ref_list), (*batch_shape, n_var, n_var)) - assert_close(res, xp.asarray(ref)) - - def test_axis_with_weights(self, xp: ArrayNamespace): - # axis=-2 (observations on first of 2D) combined with weights: - # verifies that moveaxis and weight alignment cooperate. - rng = np.random.default_rng(20260417) - m = rng.random((15, 3)) # observations on axis 0 - aw = rng.random(15) - fw = np.asarray([1, 2, 1, 3, 1, 2, 1, 1, 2, 1, 1, 1, 2, 1, 1], dtype=np.int64) - ref = np.cov(m, rowvar=False, fweights=fw, aweights=aw) - res = cov( - xp.asarray(m), - axis=-2, - fweights=xp.asarray(fw), - aweights=xp.asarray(aw), - ) - assert_close(res, xp.asarray(ref)) - - def test_axis_out_of_bounds(self, xp: ArrayNamespace): - m = xp.asarray([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - with pytest.raises(IndexError): - _ = cov(m, axis=5) - - def test_weights_wrong_ndim(self, xp: ArrayNamespace): - m = xp.asarray([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - w2d = xp.asarray([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]) - # Non-integer correction forces the generic path where the - # validation lives; native backends raise for the same reason. - with pytest.raises((ValueError, TypeError)): - _ = cov(m, correction=0.5, fweights=w2d) - with pytest.raises((ValueError, TypeError)): - _ = cov(m, correction=0.5, aweights=w2d) - - def test_weights_wrong_length(self, xp: ArrayNamespace): - m = xp.asarray([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - w_bad = xp.asarray([1.0, 1.0]) # expected length 3 - with pytest.raises((ValueError, RuntimeError)): - _ = cov(m, correction=0.5, fweights=w_bad) - with pytest.raises((ValueError, RuntimeError)): - _ = cov(m, correction=0.5, aweights=w_bad) - - def test_weights_unknown_length(self, da: ArrayNamespace): - m_np = np.asarray([[1.0, 2.0, 3.0], [4.0, 6.0, 8.0]]) - weights_np = np.asarray([1.0, 2.0, 3.0]) - keep_np = np.asarray([True, False, True]) - - keep = da.asarray(keep_np) - m = da.asarray(m_np)[:, keep] - weights = da.asarray(weights_np)[keep] - assert math.isnan(m.shape[-1]) - assert math.isnan(weights.shape[0]) - - actual = cov(m, aweights=weights) - desired = np.cov(m_np[:, keep_np], aweights=weights_np[keep_np]) - assert_close(actual, da.asarray(desired)) - - def test_weights_dof_warning_eager(self): - xp = array_namespace(cast(Array, cast(object, np.empty(0)))) - m = xp.asarray([[1.0, 2.0], [3.0, 4.0]]) - weights = xp.asarray([1.0, 1.0]) - - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - _ = cov(m, correction=2.5, aweights=weights) - assert any( - isinstance(warning.message, RuntimeWarning) - and "Degrees of freedom <= 0" in str(warning.message) - for warning in caught - ) - - def test_torch_autograd(self, torch: ArrayNamespace): - # The batched (generic) path must not detach gradients or mutate the - # input tensor in place, as `xp.asarray` does on torch. - xp = torch - rng = np.random.default_rng(20260417) - m = xp.asarray(rng.random((4, 3, 20)), dtype=xp.float64) - m.requires_grad_(True) - m_before = m.detach().clone() - # cov returns the array-api `Array` type; at runtime it is a torch - # tensor, so cast to access autograd attributes without type errors. - c = cast(Any, cov(m)) # batched -> generic path - assert c.requires_grad - assert m.requires_grad # input tensor not mutated - assert_equal(m.detach(), m_before) - c.sum().backward() - assert m.grad is not None - assert bool(xp.all(xp.isfinite(m.grad))) - - -@pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no arange", strict=False) -class TestOneHot: - @pytest.mark.parametrize("n_dim", range(4)) - @pytest.mark.parametrize("num_classes", [1, 3, 10]) - def test_dims_and_classes(self, xp: ArrayNamespace, n_dim: int, num_classes: int): - shape = tuple(range(2, 2 + n_dim)) - rng = np.random.default_rng(2347823) - np_x = rng.integers(num_classes, size=shape) - x = xp.asarray(np_x) - y = one_hot(x, num_classes) - assert y.shape == (*x.shape, num_classes) - for *i_list, j in ndindex(*shape, num_classes): - i = tuple(i_list) - assert float(y[(*i, j)]) == (int(x[i]) == j) - - def test_basic(self, xp: ArrayNamespace): - actual = one_hot(xp.asarray([0, 1, 2]), 3) - expected = xp.asarray([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]) - assert_equal(actual, expected) - - actual = one_hot(xp.asarray([1, 2, 0]), 3) - expected = xp.asarray([[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]]) - assert_equal(actual, expected) - - def test_2d(self, xp: ArrayNamespace): - actual = one_hot(xp.asarray([[2, 1, 0], [1, 0, 2]]), 3, axis=1) - expected = xp.asarray( - [ - [[0.0, 0.0, 1.0], [0.0, 1.0, 0.0], [1.0, 0.0, 0.0]], - [[0.0, 1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]], - ] - ) - assert_equal(actual, expected) - - @pytest.mark.skip_xp_backend( - Backend.ARRAY_API_STRICTEST, reason="backend doesn't support Boolean indexing" - ) - def test_abstract_size(self, xp: ArrayNamespace): - x = xp.arange(5) - x = x[x > 2] - actual = one_hot(x, 5) - expected = xp.asarray([[0.0, 0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 0.0, 1.0]]) - assert_equal(actual, expected) - - @pytest.mark.skip_xp_backend( - Backend.TORCH_GPU, reason="Puts Pytorch into a bad state." - ) - def test_out_of_bound(self, xp: ArrayNamespace): - # Undefined behavior. Either return zero, or raise. - try: - actual = one_hot(xp.asarray([-1, 3]), 3) - except IndexError: - return - expected = xp.asarray([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]) - assert_equal(actual, expected) - - @pytest.mark.parametrize( - "int_dtype", - ["int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64"], - ) - def test_int_types(self, xp: ArrayNamespace, int_dtype: str): - dtype = getattr(xp, int_dtype) - x = xp.asarray([0, 1, 2], dtype=dtype) - actual = one_hot(x, 3) - expected = xp.asarray([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]) - assert_equal(actual, expected) - - def test_custom_dtype(self, xp: ArrayNamespace): - actual = one_hot(xp.asarray([0, 1, 2], dtype=xp.int32), 3, dtype=xp.bool) - expected = xp.asarray( - [[True, False, False], [False, True, False], [False, False, True]] - ) - assert_equal(actual, expected) - - def test_axis(self, xp: ArrayNamespace): - expected = xp.asarray([[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]]).T - actual = one_hot(xp.asarray([1, 2, 0]), 3, axis=0) - assert_equal(actual, expected) - - actual = one_hot(xp.asarray([1, 2, 0]), 3, axis=-2) - assert_equal(actual, expected) - - def test_non_integer(self, xp: ArrayNamespace): - with pytest.raises(TypeError): - _ = one_hot(xp.asarray([1.0]), 3) - - def test_device(self, xp: ArrayNamespace, device: Device): - x = xp.asarray([0, 1, 2], device=device) - y = one_hot(x, 3) - assert get_device(y) == device - - -@pytest.mark.skip_xp_backend( - Backend.SPARSE, reason="read-only backend without .at support" -) -class TestCreateDiagonal: - def test_1d_from_numpy(self, xp: ArrayNamespace): - # from np.diag tests - vals = 100 * xp.arange(5, dtype=xp.float64) - b = xp.zeros((5, 5), dtype=xp.float64) - for k in range(5): - b = at(b)[k, k].set(vals[k]) - assert_equal(create_diagonal(vals), b) - b = xp.zeros((7, 7), dtype=xp.float64) - c = xp.asarray(b, copy=True) - for k in range(5): - b = at(b)[k, k + 2].set(vals[k]) - c = at(c)[k + 2, k].set(vals[k]) - assert_equal(create_diagonal(vals, offset=2), b) - assert_equal(create_diagonal(vals, offset=-2), c) - - @pytest.mark.parametrize("n", range(1, 10)) - @pytest.mark.parametrize("offset", range(1, 10)) - def test_1d_from_scipy(self, xp: ArrayNamespace, n: int, offset: int): - # from scipy._lib tests - rng = np.random.default_rng(2347823) - one = xp.asarray(1.0) - x = rng.random(n) - A = create_diagonal(xp.asarray(x, dtype=one.dtype), offset=offset) - B = xp.asarray(np.diag(x, offset), dtype=one.dtype) - assert_equal(A, B) - - def test_0d_raises(self, xp: ArrayNamespace): - with pytest.raises(ValueError, match="1-dimensional"): - _ = create_diagonal(xp.asarray(1)) - - @pytest.mark.parametrize( - "shape", - [ - (0,), - (10,), - (0, 1), - (1, 0), - (0, 0), - (2, 3), - (4, 2, 1), - (1, 1, 7), - (0, 0, 1), - (3, 2, 4, 5), - ], - ) - def test_nd(self, xp: ArrayNamespace, shape: tuple[int, ...]): - rng = np.random.default_rng(2347823) - b = xp.asarray( - rng.integers((1 << 64) - 1, size=shape, dtype=np.uint64), dtype=xp.uint64 - ) - c = create_diagonal(b) - zero = xp.zeros((), dtype=xp.uint64) - assert c.shape == (*b.shape, b.shape[-1]) - for i in ndindex(*eager_shape(c)): - assert_equal(c[i], b[i[:-1]] if i[-2] == i[-1] else zero) - - def test_device(self, xp: ArrayNamespace, device: Device): - x = xp.asarray([1, 2, 3], device=device) - assert get_device(create_diagonal(x)) == device - - def test_xp(self, xp: ArrayNamespace): - x = xp.asarray([1, 2]) - y = create_diagonal(x, xp=xp) - assert_equal(y, xp.asarray([[1, 0], [0, 2]])) - - -class TestDefaultDType: - def test_basic(self, xp: ArrayNamespace): - assert default_dtype(xp) == xp.empty(0).dtype - - def test_kind(self, xp: ArrayNamespace): - assert default_dtype(xp, "real floating") == xp.empty(0).dtype - assert default_dtype(xp, "complex floating") == (xp.empty(0) * 1j).dtype - assert default_dtype(xp, "integral") == xp.int64 - assert default_dtype(xp, "indexing") == xp.int64 - - with pytest.raises(ValueError, match="Unknown kind"): - _ = default_dtype(xp, "foo") # type: ignore[arg-type] # pyright: ignore[reportArgumentType] - - def test_device(self, xp: ArrayNamespace, device: Device): - # Note: at the moment there are no known namespaces with - # device-specific default dtypes. - assert default_dtype(xp, device=None) == xp.empty(0).dtype - assert default_dtype(xp, device=device) == xp.empty(0).dtype - - def test_torch(self, torch: ArrayNamespace): - xp = torch - xp.set_default_dtype(xp.float64) - assert default_dtype(xp) == xp.float64 - assert default_dtype(xp, "real floating") == xp.float64 - assert default_dtype(xp, "complex floating") == xp.complex128 - - xp.set_default_dtype(xp.float32) - assert default_dtype(xp) == xp.float32 - assert default_dtype(xp, "real floating") == xp.float32 - assert default_dtype(xp, "complex floating") == xp.complex64 - - -@pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no arange", strict=False) -class TestDiagIndices: - def test_basic(self, xp: ArrayNamespace): - rows, cols = diag_indices(5, xp=xp) - ref_rows, ref_cols = np.diag_indices(5) - assert_equal(rows, xp.asarray(ref_rows)) - assert_equal(cols, xp.asarray(ref_cols)) - - @pytest.mark.parametrize("n", [2, 4, 7]) - @pytest.mark.parametrize("ndim", [1, 2, 3, 4]) - def test_ndim(self, xp: ArrayNamespace, n: int, ndim: int): - idx = diag_indices(n, ndim=ndim, xp=xp) - assert len(idx) == ndim - ref = np.diag_indices(n, ndim=ndim) - for got, expected in zip(idx, ref, strict=True): - assert_equal(got, xp.asarray(expected)) - - def test_empty(self, xp: ArrayNamespace): - rows, cols = diag_indices(0, xp=xp) - assert rows.shape == (0,) - assert cols.shape == (0,) - - def test_validation(self, xp: ArrayNamespace): - with pytest.raises(ValueError, match="`n` must be non-negative"): - _ = diag_indices(-1, xp=xp) - with pytest.raises(ValueError, match="`ndim` must be >= 1"): - _ = diag_indices(3, ndim=0, xp=xp) - - def test_device(self, xp: ArrayNamespace, device: Device): - default_device = get_device(xp.empty(0)) - rows, cols = diag_indices(3, device=None, xp=xp) - assert get_device(rows) == default_device - assert get_device(cols) == default_device - rows, cols = diag_indices(3, device=device, xp=xp) - assert get_device(rows) == device - assert get_device(cols) == device - - -@pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no arange/nonzero", strict=False) -@pytest.mark.xfail_xp_backend( - Backend.ARRAY_API_STRICTEST, - reason="generic path uses nonzero (data-dependent)", - strict=False, -) -@pytest.mark.parametrize( - ("xpx_fn", "np_fn"), - [(tril_indices, np.tril_indices), (triu_indices, np.triu_indices)], - ids=["tril", "triu"], -) -class TestTriIndices: - def test_basic( - self, - xp: ArrayNamespace, - xpx_fn: Callable[..., tuple[Array, Array]], - np_fn: Callable[..., tuple[Array, Array]], - ): - rows, cols = xpx_fn(4, xp=xp) - ref_rows, ref_cols = np_fn(4) - assert_equal(rows, xp.asarray(ref_rows)) - assert_equal(cols, xp.asarray(ref_cols)) - - @pytest.mark.parametrize("offset", [-2, -1, 0, 1, 2]) - def test_offset( - self, - xp: ArrayNamespace, - xpx_fn: Callable[..., tuple[Array, Array]], - np_fn: Callable[..., tuple[Array, Array]], - offset: int, - ): - rows, cols = xpx_fn(5, offset=offset, xp=xp) - ref_rows, ref_cols = np_fn(5, k=offset) - assert_equal(rows, xp.asarray(ref_rows)) - assert_equal(cols, xp.asarray(ref_cols)) - - def test_rectangular( - self, - xp: ArrayNamespace, - xpx_fn: Callable[..., tuple[Array, Array]], - np_fn: Callable[..., tuple[Array, Array]], - ): - rows, cols = xpx_fn(3, m=5, xp=xp) - ref_rows, ref_cols = np_fn(3, m=5) - assert_equal(rows, xp.asarray(ref_rows)) - assert_equal(cols, xp.asarray(ref_cols)) - - @pytest.mark.xfail_xp_backend( - Backend.DASK, reason="dask: no 2D fancy indexing", strict=False - ) - def test_use_to_read( - self, - xp: ArrayNamespace, - xpx_fn: Callable[..., tuple[Array, Array]], - np_fn: Callable[..., tuple[Array, Array]], - ): - rng = np.random.default_rng(0) - a = rng.integers(0, 100, (4, 4)) - a_xp = xp.asarray(a) - rows, cols = xpx_fn(4, xp=xp) - assert_equal(a_xp[rows, cols], xp.asarray(a[np_fn(4)])) - - def test_validation( - self, - xp: ArrayNamespace, - xpx_fn: Callable[..., tuple[Array, Array]], - np_fn: Callable[..., tuple[Array, Array]], # noqa: ARG002 # pytest param - ): - with pytest.raises(ValueError, match="`n` must be non-negative"): - _ = xpx_fn(-1, xp=xp) - with pytest.raises(ValueError, match="`m` must be non-negative"): - _ = xpx_fn(3, m=-1, xp=xp) - - def test_device( - self, - xp: ArrayNamespace, - device: Device, - xpx_fn: Callable[..., tuple[Array, Array]], - np_fn: Callable[..., tuple[Array, Array]], # noqa: ARG002 # pytest param - ): - default_device = get_device(xp.empty(0)) - rows, cols = xpx_fn(4, device=None, xp=xp) - assert get_device(rows) == default_device - assert get_device(cols) == default_device - rows, cols = xpx_fn(4, device=device, xp=xp) - assert get_device(rows) == device - assert get_device(cols) == device - - -@pytest.mark.filterwarnings(r"ignore:.*removed in v1.0.0.*:DeprecationWarning") -class TestExpandDims: - def test_single_axis(self, xp: ArrayNamespace): - """Trivial case where xpx.expand_dims doesn't add anything to xp.expand_dims""" - a = xp.asarray(np.reshape(np.arange(2 * 3 * 4 * 5), (2, 3, 4, 5))) - for axis in range(-5, 4): - b = expand_dims(a, axis=axis) - assert_equal(b, xp.expand_dims(a, axis=axis)) - - def test_axis_tuple(self, xp: ArrayNamespace): - a = xp.empty((3, 3, 3)) - assert expand_dims(a, axis=(0, 1, 2)).shape == (1, 1, 1, 3, 3, 3) - assert expand_dims(a, axis=(0, -1, -2)).shape == (1, 3, 3, 3, 1, 1) - assert expand_dims(a, axis=(0, 3, 5)).shape == (1, 3, 3, 1, 3, 1) - assert expand_dims(a, axis=(0, -3, -5)).shape == (1, 1, 3, 1, 3, 3) - - def test_axis_out_of_range(self, xp: ArrayNamespace): - a = xp.empty((2, 3, 4, 5)) - with pytest.raises(IndexError, match="out of bounds"): - _ = expand_dims(a, axis=-6) - with pytest.raises(IndexError, match="out of bounds"): - _ = expand_dims(a, axis=5) - - a = xp.empty((3, 3, 3)) - with pytest.raises(IndexError, match="out of bounds"): - _ = expand_dims(a, axis=(0, -6)) - with pytest.raises(IndexError, match="out of bounds"): - _ = expand_dims(a, axis=(0, 5)) - - def test_repeated_axis(self, xp: ArrayNamespace): - a = xp.empty((3, 3, 3)) - with pytest.raises(ValueError, match="Duplicate dimensions"): - _ = expand_dims(a, axis=(1, 1)) - - def test_positive_negative_repeated(self, xp: ArrayNamespace): - # https://github.com/data-apis/array-api/issues/760#issuecomment-1989449817 - a = xp.empty((2, 3, 4, 5)) - with pytest.raises(ValueError, match="Duplicate dimensions"): - _ = expand_dims(a, axis=(3, -3)) - - def test_device(self, xp: ArrayNamespace, device: Device): - x = xp.asarray([1, 2, 3], device=device) - assert get_device(expand_dims(x, axis=0)) == device - - def test_xp(self, xp: ArrayNamespace): - x = xp.asarray([1, 2, 3]) - y = expand_dims(x, axis=(0, 1, 2), xp=xp) - assert y.shape == (1, 1, 1, 3) - - -@pytest.mark.filterwarnings( # array_api_strictest - "ignore:invalid value encountered:RuntimeWarning:array_api_strict" -) -@pytest.mark.filterwarnings( # sparse - "ignore:invalid value encountered:RuntimeWarning:sparse" -) -class TestIsClose: - @pytest.mark.parametrize("swap", [False, True]) - @pytest.mark.parametrize( - ("a", "b"), - [ - (0.0, 0.0), - (1.0, 1.0), - (1.0, 2.0), - (1.0, -1.0), - (100.0, 101.0), - (0, 0), - (1, 1), - (1, 2), - (1, -1), - (1.0 + 1j, 1.0 + 1j), - (1.0 + 1j, 1.0 - 1j), - (float("inf"), float("inf")), - (float("inf"), 100.0), - (float("inf"), float("-inf")), - (float("-inf"), float("-inf")), - (float("nan"), float("nan")), - (float("nan"), 100.0), - (1e6, 1e6 + 1), # True - within rtol - (1e6, 1e6 + 100), # False - outside rtol - (1e-6, 1.1e-6), # False - outside atol - (1e-7, 1.1e-7), # True - outside atol - (1e6 + 0j, 1e6 + 1j), # True - within rtol - (1e6 + 0j, 1e6 + 100j), # False - outside rtol - ], - ) - def test_basic(self, a: float, b: float, swap: bool, xp: ArrayNamespace): - if swap: - b, a = a, b - a_xp = xp.asarray(a) - b_xp = xp.asarray(b) - - assert_equal(isclose(a_xp, b_xp), xp.asarray(np.isclose(a, b))) - - with warnings.catch_warnings(): - warnings.simplefilter("ignore") - ar_np = a * np.arange(10) - br_np = b * np.arange(10) - ar_xp = xp.asarray(ar_np) - br_xp = xp.asarray(br_np) - - assert_equal(isclose(ar_xp, br_xp), xp.asarray(np.isclose(ar_np, br_np))) - - @pytest.mark.parametrize("dtype", ["float32", "int32"]) - def test_broadcast(self, dtype: str, xp: ArrayNamespace): - dtype = getattr(xp, dtype) - a = xp.asarray([1, 2, 3], dtype=dtype) - b = xp.asarray([[1], [5]], dtype=dtype) - actual = isclose(a, b) - expect = xp.asarray( - [[True, False, False], [False, False, False]], dtype=xp.bool - ) - - assert_equal(actual, expect) - - def test_some_inf(self, xp: ArrayNamespace): - a = xp.asarray([0.0, 1.0, xp.inf, xp.inf, xp.inf]) - b = xp.asarray([1e-9, 1.0, xp.inf, -xp.inf, 2.0]) - actual = isclose(a, b) - assert_equal(actual, xp.asarray([True, True, True, False, False])) - - def test_equal_nan(self, xp: ArrayNamespace): - a = xp.asarray([xp.nan, xp.nan, 1.0]) - b = xp.asarray([xp.nan, 1.0, xp.nan]) - assert_equal(isclose(a, b), xp.asarray([False, False, False])) - assert_equal(isclose(a, b, equal_nan=True), xp.asarray([True, False, False])) - - @pytest.mark.parametrize("dtype", ["float32", "complex64", "int32"]) - def test_tolerance(self, dtype: str, xp: ArrayNamespace): - dtype = getattr(xp, dtype) - a = xp.asarray([100, 100], dtype=dtype) - b = xp.asarray([101, 102], dtype=dtype) - assert_equal(isclose(a, b), xp.asarray([False, False])) - assert_equal(isclose(a, b, atol=1), xp.asarray([True, False])) - assert_equal(isclose(a, b, rtol=0.01), xp.asarray([True, False])) - - # Attempt to trigger division by 0 in rtol on int dtype - assert_equal(isclose(a, b, rtol=0), xp.asarray([False, False])) - assert_equal(isclose(a, b, atol=1, rtol=0), xp.asarray([True, False])) - - @pytest.mark.parametrize("dtype", ["int8", "uint8"]) - def test_tolerance_integer_overflow(self, dtype: str, xp: ArrayNamespace): - """1/rtol is too large for dtype""" - a = xp.asarray([100, 100], dtype=getattr(xp, dtype)) - b = xp.asarray([100, 101], dtype=getattr(xp, dtype)) - assert_equal(isclose(a, b), xp.asarray([True, False])) - - def test_very_small_numbers(self, xp: ArrayNamespace): - a = xp.asarray([1e-9, 1e-9]) - b = xp.asarray([1.0001e-9, 1.00001e-9]) - # Difference is below default atol - assert_equal(isclose(a, b), xp.asarray([True, True])) - # Use only rtol - assert_equal(isclose(a, b, atol=0), xp.asarray([False, True])) - assert_equal(isclose(a, b, atol=0, rtol=0), xp.asarray([False, False])) - - def test_bool_dtype(self, xp: ArrayNamespace): - a = xp.asarray([False, True, False]) - b = xp.asarray([True, True, False]) - assert_equal(isclose(a, b), xp.asarray([False, True, True])) - assert_equal(isclose(a, b, atol=1), xp.asarray([True, True, True])) - assert_equal(isclose(a, b, atol=2), xp.asarray([True, True, True])) - assert_equal(isclose(a, b, rtol=1), xp.asarray([True, True, True])) - assert_equal(isclose(a, b, rtol=2), xp.asarray([True, True, True])) - - # Test broadcasting - assert_equal( - isclose(a, xp.asarray(True), atol=1), xp.asarray([True, True, True]) - ) - assert_equal( - isclose(xp.asarray(True), b, atol=1), xp.asarray([True, True, True]) - ) - - @pytest.mark.skip_xp_backend(Backend.SPARSE, reason="index by sparse array") - @pytest.mark.skip_xp_backend(Backend.ARRAY_API_STRICTEST, reason="unknown shape") - def test_none_shape(self, xp: ArrayNamespace): - a = xp.asarray([1, 5, 0]) - b = xp.asarray([1, 4, 2]) - b = b[a < 5] - a = a[a < 5] - assert_equal(isclose(a, b), xp.asarray([True, False])) - - @pytest.mark.skip_xp_backend(Backend.SPARSE, reason="index by sparse array") - @pytest.mark.skip_xp_backend(Backend.ARRAY_API_STRICTEST, reason="unknown shape") - def test_none_shape_bool(self, xp: ArrayNamespace): - a = xp.asarray([True, True, False]) - b = xp.asarray([True, False, True]) - b = b[a] - a = a[a] - assert_equal(isclose(a, b), xp.asarray([True, False])) - - @pytest.mark.skip_xp_backend(Backend.NUMPY_READONLY, reason="xp=xp") - def test_python_scalar(self, xp: ArrayNamespace): - a = xp.asarray([0.0, 0.1], dtype=xp.float32) - assert_equal(isclose(a, 0.0), xp.asarray([True, False])) - assert_equal(isclose(0.0, a), xp.asarray([True, False])) - - a = xp.asarray([0, 1], dtype=xp.int16) - assert_equal(isclose(a, 0), xp.asarray([True, False])) - assert_equal(isclose(0, a), xp.asarray([True, False])) - - assert_equal(isclose(0, 0, xp=xp), xp.asarray(True)) - assert_equal(isclose(0, 1, xp=xp), xp.asarray(False)) - - def test_all_python_scalars(self): - with pytest.raises(TypeError, match=r"array_namespace requires .* array input"): - _ = isclose(0, 0) - - def test_xp(self, xp: ArrayNamespace): - a = xp.asarray([0.0, 0.0]) - b = xp.asarray([1e-9, 1e-4]) - assert_equal(isclose(a, b, xp=xp), xp.asarray([True, False])) - - @pytest.mark.parametrize("equal_nan", [True, False]) - def test_device(self, xp: ArrayNamespace, device: Device, equal_nan: bool): - a = xp.asarray([0.0, 0.0, xp.nan], device=device) - b = xp.asarray([1e-9, 1e-4, xp.nan], device=device) - res = isclose(a, b, equal_nan=equal_nan) - assert get_device(res) == device - - def test_array_on_device_with_scalar(self, xp: ArrayNamespace, device: Device): - a = xp.asarray([0.01, 0.5, 0.8, 0.9, 1.00001], device=device, dtype=xp.float64) - b = 1 - res = isclose(a, b) - assert get_device(res) == device - assert_equal(res, xp.asarray([False, False, False, False, True], device=device)) - - a = 0.1 - b = xp.asarray([0.01, 0.5, 0.8, 0.9, 0.100001], device=device, dtype=xp.float64) - res = isclose(a, b) - assert get_device(res) == device - assert_equal(res, xp.asarray([False, False, False, False, True], device=device)) - - -class TestKron: - def test_basic(self, xp: ArrayNamespace): - # Using 0-dimensional array - a = xp.asarray(1) - b = xp.asarray([[1, 2], [3, 4]]) - assert_equal(kron(a, b), b) - assert_equal(kron(b, a), b) - - # Using 1-dimensional array - a = xp.asarray([3]) - b = xp.asarray([[1, 2], [3, 4]]) - k = xp.asarray([[3, 6], [9, 12]]) - assert_equal(kron(a, b), k) - assert_equal(kron(b, a), k) - - # Using 3-dimensional array - a = xp.asarray([[[1]], [[2]]]) - b = xp.asarray([[1, 2], [3, 4]]) - k = xp.asarray([[[1, 2], [3, 4]], [[2, 4], [6, 8]]]) - assert_equal(kron(a, b), k) - assert_equal(kron(b, a), k) - - def test_kron_smoke(self, xp: ArrayNamespace): - a = xp.ones((3, 3)) - b = xp.ones((3, 3)) - k = xp.ones((9, 9)) - - assert_equal(kron(a, b), k) - - @pytest.mark.parametrize( - ("shape_a", "shape_b"), - [ - ((1, 1), (1, 1)), - ((1, 2, 3), (4, 5, 6)), - ((2, 2), (2, 2, 2)), - ((1, 0), (1, 1)), - ((2, 0, 2), (2, 2)), - ((2, 0, 0, 2), (2, 0, 2)), - ], - ) - def test_kron_shape( - self, xp: ArrayNamespace, shape_a: tuple[int, ...], shape_b: tuple[int, ...] - ): - a = xp.ones(shape_a) - b = xp.ones(shape_b) - normalised_shape_a = xp.asarray( - (1,) * max(0, len(shape_b) - len(shape_a)) + shape_a - ) - normalised_shape_b = xp.asarray( - (1,) * max(0, len(shape_a) - len(shape_b)) + shape_b - ) - expected_shape = tuple( - int(dim) for dim in xp.multiply(normalised_shape_a, normalised_shape_b) - ) - - k = kron(a, b) - assert k.shape == expected_shape - - def test_python_scalar(self, xp: ArrayNamespace): - a = 1 - # Test no dtype promotion to xp.asarray(a); use b.dtype - b = xp.asarray([[1, 2], [3, 4]], dtype=xp.int16) - assert_equal(kron(a, b), b) - assert_equal(kron(b, a), b) - assert_equal(kron(1, 1, xp=xp), xp.asarray(1)) - - def test_all_python_scalars(self): - with pytest.raises(TypeError, match=r"array_namespace requires .* array input"): - _ = kron(1, 1) - - def test_device(self, xp: ArrayNamespace, device: Device): - x1 = xp.asarray([1, 2, 3], device=device) - x2 = xp.asarray([4, 5], device=device) - assert get_device(kron(x1, x2)) == device - - def test_xp(self, xp: ArrayNamespace): - a = xp.ones((3, 3)) - b = xp.ones((3, 3)) - k = xp.ones((9, 9)) - assert_equal(kron(a, b, xp=xp), k) - - -class TestNanToNum: - def test_bool(self, xp: ArrayNamespace) -> None: - a = xp.asarray([True]) - assert_equal(nan_to_num(a, xp=xp), a) - - def test_scalar_pos_inf(self, xp: ArrayNamespace, infinity: float) -> None: - a = xp.inf - assert_equal(nan_to_num(a, xp=xp), xp.asarray(infinity)) - - def test_scalar_neg_inf(self, xp: ArrayNamespace, infinity: float) -> None: - a = -xp.inf - assert_equal(nan_to_num(a, xp=xp), -xp.asarray(infinity)) - - def test_scalar_nan(self, xp: ArrayNamespace) -> None: - a = xp.nan - assert_equal(nan_to_num(a, xp=xp), xp.asarray(0.0)) - - def test_real(self, xp: ArrayNamespace, infinity: float) -> None: - a = xp.asarray([xp.inf, -xp.inf, xp.nan, -128, 128]) - assert_equal( - nan_to_num(a, xp=xp), - xp.asarray( - [ - infinity, - -infinity, - 0.0, - -128, - 128, - ] - ), - ) - - def test_complex(self, xp: ArrayNamespace, infinity: float) -> None: - a = xp.asarray( - [ - complex(xp.inf, xp.nan), - xp.nan, - complex(xp.nan, xp.inf), - ] - ) - assert_equal( - nan_to_num(a), - xp.asarray([complex(infinity, 0), complex(0, 0), complex(0, infinity)]), - ) - - def test_empty_array(self, xp: ArrayNamespace) -> None: - a = xp.asarray([], dtype=xp.float32) # forced dtype due to torch - assert_equal(nan_to_num(a, xp=xp), a) - assert xp.isdtype(nan_to_num(a, xp=xp).dtype, xp.float32) - - @pytest.mark.parametrize( - ("in_vals", "fill_value", "out_vals"), - [ - ([1, 2, np.nan, 4], 3, [1.0, 2.0, 3.0, 4.0]), - ([1, 2, np.nan, 4], 3.0, [1.0, 2.0, 3.0, 4.0]), - ( - [ - complex(1, 1), - complex(2, 2), - complex(np.nan, 0), - complex(4, 4), - ], - 3, - [ - complex(1.0, 1.0), - complex(2.0, 2.0), - complex(3.0, 0.0), - complex(4.0, 4.0), - ], - ), - ( - [ - complex(1, 1), - complex(2, 2), - complex(0, np.nan), - complex(4, 4), - ], - 3.0, - [ - complex(1.0, 1.0), - complex(2.0, 2.0), - complex(0.0, 3.0), - complex(4.0, 4.0), - ], - ), - ( - [ - complex(1, 1), - complex(2, 2), - complex(np.nan, np.nan), - complex(4, 4), - ], - 3.0, - [ - complex(1.0, 1.0), - complex(2.0, 2.0), - complex(3.0, 3.0), - complex(4.0, 4.0), - ], - ), - ], - ) - def test_fill_value_success( - self, - xp: ArrayNamespace, - in_vals: Array, - fill_value: float, - out_vals: Array, - ) -> None: - a = xp.asarray(in_vals) - assert_equal( - nan_to_num(a, fill_value=fill_value, xp=xp), - xp.asarray(out_vals), - ) - - def test_fill_value_failure(self, xp: ArrayNamespace) -> None: - a = xp.asarray( - [ - complex(1, 1), - complex(xp.nan, xp.nan), - complex(3, 3), - ] - ) - with pytest.raises( - TypeError, - match="Complex fill values are not supported", - ): - _ = nan_to_num( - a, - fill_value=complex(2, 2), # type: ignore[arg-type] # pyright: ignore[reportArgumentType] - xp=xp, - ) - - -class TestNUnique: - @pytest.mark.skip_xp_backend( - Backend.ARRAY_API_STRICT, reason="array-agnostic fallback" - ) - @pytest.mark.skip_xp_backend( - Backend.ARRAY_API_STRICTEST, reason="array-agnostic fallback" - ) - @pytest.mark.skip_xp_backend(Backend.DASK, reason="array-agnostic fallback") - @pytest.mark.skip_xp_backend(Backend.SPARSE, reason="array-agnostic fallback") - def test_delegates( - self, - xp: ArrayNamespace, - monkeypatch: pytest.MonkeyPatch, - ): - def fallback(*_args: object, **_kwargs: object) -> Array: - msg = "array-agnostic fallback should not be used" - raise AssertionError(msg) - - monkeypatch.setattr(functions, "nunique", fallback) - a = xp.asarray([1, 1, 2]) - assert_equal(nunique(a), xp.asarray(2)) - - def test_simple(self, xp: ArrayNamespace): - a = xp.asarray([[1, 1], [0, 2], [2, 2]]) - assert_equal(nunique(a), xp.asarray(3)) - - def test_empty(self, xp: ArrayNamespace): - a = xp.asarray([]) - assert_equal(nunique(a), xp.asarray(0)) - - def test_size1(self, xp: ArrayNamespace): - a = xp.asarray([123]) - assert_equal(nunique(a), xp.asarray(1)) - - def test_all_equal(self, xp: ArrayNamespace): - a = xp.asarray([123, 123, 123]) - assert_equal(nunique(a), xp.asarray(1)) - - @pytest.mark.xfail_xp_backend(Backend.DASK, reason="No equal_nan kwarg in unique") - def test_nan(self, xp: ArrayNamespace, library: Backend): - if library.like(Backend.NUMPY) and NUMPY_VERSION < (1, 24): - pytest.xfail("NumPy <1.24 has no equal_nan kwarg in unique") - - # Each NaN is counted separately - a = xp.asarray([xp.nan, 123.0, xp.nan]) - assert_equal(nunique(a), xp.asarray(3)) - - @pytest.mark.parametrize("size", [0, 1, 2]) - def test_device(self, xp: ArrayNamespace, device: Device, size: int): - a = xp.asarray([0.0] * size, device=device) - assert get_device(nunique(a)) == device - - def test_xp(self, xp: ArrayNamespace): - a = xp.asarray([[1, 1], [0, 2], [2, 2]]) - assert_equal(nunique(a, xp=xp), xp.asarray(3)) - - -class TestPad: - def test_simple(self, xp: ArrayNamespace): - a = xp.asarray([1, 2, 3]) - padded = pad(a, 2) - assert_equal(padded, xp.asarray([0, 0, 1, 2, 3, 0, 0])) - - @pytest.mark.xfail_xp_backend( - Backend.SPARSE, reason="constant_values can only be equal to fill value" - ) - def test_fill_value(self, xp: ArrayNamespace): - a = xp.asarray([1, 2, 3]) - padded = pad(a, 2, constant_values=42) - assert_equal(padded, xp.asarray([42, 42, 1, 2, 3, 42, 42])) - - def test_ndim(self, xp: ArrayNamespace): - a = xp.asarray(np.reshape(np.arange(2 * 3 * 4), (2, 3, 4))) - padded = pad(a, 2) - assert padded.shape == (6, 7, 8) - - def test_mode_not_implemented(self, xp: ArrayNamespace): - a = xp.asarray([1, 2, 3]) - with pytest.raises(NotImplementedError, match="Only `'constant'`"): - _ = pad(a, 2, mode="edge") # type: ignore[arg-type] # pyright: ignore[reportArgumentType] - - def test_device(self, xp: ArrayNamespace, device: Device): - a = xp.asarray(0.0, device=device) - assert get_device(pad(a, 2)) == device - - def test_xp(self, xp: ArrayNamespace): - padded = pad(xp.asarray(0), 1, xp=xp) - assert_equal(padded, xp.asarray(0)) - - def test_tuple_width(self, xp: ArrayNamespace): - a = xp.asarray(np.reshape(np.arange(12), (3, 4))) - padded = pad(a, (1, 0)) - assert padded.shape == (4, 5) - - padded = pad(a, (1, 2)) - assert padded.shape == (6, 7) - - with pytest.raises((ValueError, RuntimeError)): - _ = pad(a, [(1, 2, 3)]) # type: ignore[list-item] # pyright: ignore[reportArgumentType] - - def test_sequence_of_tuples_width(self, xp: ArrayNamespace): - a = xp.asarray(np.reshape(np.arange(12), (3, 4))) - - padded = pad(a, ((1, 0), (0, 2))) - assert padded.shape == (4, 6) - padded = pad(a, ((1, 0), (0, 0))) - assert padded.shape == (4, 4) - - -assume_unique = pytest.mark.parametrize( - "assume_unique", - [ - True, - pytest.param( - False, - marks=pytest.mark.xfail_xp_backend( - Backend.DASK, reason="NaN-shaped arrays" - ), - ), - ], -) - - -@pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no argsort") -@pytest.mark.skip_xp_backend(Backend.ARRAY_API_STRICTEST, reason="no unique_values") -class TestSetDiff1D: - @pytest.mark.xfail_xp_backend(Backend.DASK, reason="NaN-shaped arrays") - @pytest.mark.xfail_xp_backend( - Backend.TORCH, reason="index_select not implemented for uint32" - ) - @pytest.mark.xfail_xp_backend( - Backend.TORCH_GPU, reason="index_select not implemented for uint32" - ) - def test_setdiff1d(self, xp: ArrayNamespace): - x1 = xp.asarray([6, 5, 4, 7, 1, 2, 7, 4]) - x2 = xp.asarray([2, 4, 3, 3, 2, 1, 5]) - - expected = xp.asarray([6, 7]) - actual = setdiff1d(x1, x2) - assert_equal(actual, expected) - - x1 = xp.arange(21) - x2 = xp.arange(19) - expected = xp.asarray([19, 20]) - actual = setdiff1d(x1, x2) - assert_equal(actual, expected) - - assert_equal(setdiff1d(xp.empty(0), xp.empty(0)), xp.empty(0)) - x1 = xp.empty(0, dtype=xp.uint32) - x2 = x1 - assert xp.isdtype(setdiff1d(x1, x2).dtype, xp.uint32) - - def test_assume_unique(self, xp: ArrayNamespace): - x1 = xp.asarray([3, 2, 1]) - x2 = xp.asarray([7, 5, 2]) - expected = xp.asarray([3, 1]) - actual = setdiff1d(x1, x2, assume_unique=True) - assert_equal(actual, expected) - - @assume_unique - @pytest.mark.parametrize("shape1", [(), (1,), (1, 1)]) - @pytest.mark.parametrize("shape2", [(), (1,), (1, 1)]) - def test_shapes( - self, - assume_unique: bool, - shape1: tuple[int, ...], - shape2: tuple[int, ...], - xp: ArrayNamespace, - ): - x1 = xp.zeros(shape1) - x2 = xp.zeros(shape2) - - actual = setdiff1d(x1, x2, assume_unique=assume_unique) - assert_equal(actual, xp.empty((0,))) - - @assume_unique - @pytest.mark.skip_xp_backend(Backend.NUMPY_READONLY, reason="xp=xp") - def test_python_scalar(self, xp: ArrayNamespace, assume_unique: bool): - # Test no dtype promotion to xp.asarray(x2); use x1.dtype - x1 = xp.asarray([3, 1, 2], dtype=xp.int16) - x2 = 3 - actual = setdiff1d(x1, x2, assume_unique=assume_unique) - assert_equal(actual, xp.asarray([1, 2], dtype=xp.int16)) - - actual = setdiff1d(x2, x1, assume_unique=assume_unique) - assert_equal(actual, xp.asarray([], dtype=xp.int16)) - - assert_equal( - setdiff1d(0, 0, assume_unique=assume_unique, xp=xp), - xp.asarray([0])[:0], # Default int dtype for backend - ) - - @pytest.mark.parametrize("assume_unique", [True, False]) - def test_all_python_scalars(self, assume_unique: bool): - with pytest.raises(TypeError, match=r"array_namespace requires .* array input"): - _ = setdiff1d(0, 0, assume_unique=assume_unique) - - @assume_unique - @pytest.mark.skip_xp_backend( - Backend.TORCH, reason="device='meta' does not support unknown shapes" - ) - def test_device(self, xp: ArrayNamespace, device: Device, assume_unique: bool): - x1 = xp.asarray([3, 8, 20], device=device) - x2 = xp.asarray([2, 3, 4], device=device) - assert get_device(setdiff1d(x1, x2, assume_unique=assume_unique)) == device - - @pytest.mark.skip_xp_backend(Backend.NUMPY_READONLY, reason="xp=xp") - def test_xp(self, xp: ArrayNamespace): - x1 = xp.asarray([3, 8, 20]) - x2 = xp.asarray([2, 3, 4]) - expected = xp.asarray([8, 20]) - actual = setdiff1d(x1, x2, assume_unique=True, xp=xp) - assert_equal(actual, expected) - - -class TestSinc: - def test_simple(self, xp: ArrayNamespace): - assert_equal(sinc(xp.asarray(0.0)), xp.asarray(1.0)) - x = xp.asarray(np.linspace(-1, 1, 100)) - w = sinc(x) - # check symmetry - assert_close(w, xp.flip(w, axis=0)) - - @pytest.mark.parametrize("x", [0, 1 + 3j]) - def test_dtype(self, xp: ArrayNamespace, x: complex): - with pytest.raises(ValueError, match="real floating data type"): - _ = sinc(xp.asarray(x)) - - def test_3d(self, xp: ArrayNamespace): - x = np.arange(18, dtype=np.float64).reshape((3, 3, 2)) - expected = np.zeros_like(x) - expected[0, 0, 0] = 1 - x = xp.asarray(x) - expected = xp.asarray(expected) - assert_close(sinc(x), expected, atol=1e-15) - - def test_device(self, xp: ArrayNamespace, device: Device): - x = xp.asarray(0.0, device=device) - assert get_device(sinc(x)) == device - - def test_xp(self, xp: ArrayNamespace): - assert_equal(sinc(xp.asarray(0.0), xp=xp), xp.asarray(1.0)) - - -class TestPartition: - @classmethod - def _assert_valid_partition( - cls, - x_np: np.ndarray | None, - k: int, - y: Array, - xp: ArrayNamespace, - axis: int | None = -1, - ): - """ - x_np : input array - k : int - y : output array returned by the partition function to test - """ - if x_np is not None: - assert y.shape == np.partition(x_np, k, axis=axis).shape - if y.ndim != 1 and axis == 0: - assert isinstance(y.shape[1], int) - for i in range(y.shape[1]): - cls._assert_valid_partition(None, k, y[:, i, ...], xp, axis=0) - elif y.ndim != 1: - assert axis is not None - axis = axis - 1 if axis != -1 else -1 - assert isinstance(y.shape[0], int) - for i in range(y.shape[0]): - cls._assert_valid_partition(None, k, y[i, ...], xp, axis=axis) - else: - if k > 0: - assert xp.max(y[:k]) <= y[k] - assert y[k] <= xp.min(y[k:]) - - @classmethod - def _partition( - cls, x: np.ndarray, k: int, xp: ArrayNamespace, axis: int | None = -1 - ): - return partition(xp.asarray(x), k, axis=axis) - - def _test_1d(self, xp: ArrayNamespace): - rng = np.random.default_rng() - for n in [2, 3, 4, 5, 7, 10, 20, 50, 100, 1_000]: - k = int(rng.integers(n)) - x1 = rng.integers(n, size=n) - y = self._partition(x1, k, xp) - self._assert_valid_partition(x1, k, y, xp) - x2 = rng.random(n) - y = self._partition(x2, k, xp) - self._assert_valid_partition(x2, k, y, xp) - - def _test_nd(self, xp: ArrayNamespace, ndim: int): - rng = np.random.default_rng() - - for n in [2, 3, 5, 10, 20, 100]: - base_shape = [int(v) for v in rng.integers(1, 4, size=ndim)] - k = int(rng.integers(n)) - - for i in range(ndim): - shape = base_shape[:] - shape[i] = n - x = rng.integers(n, size=tuple(shape)) - y = self._partition(x, k, xp, axis=i) - self._assert_valid_partition(x, k, y, xp, axis=i) - - z = rng.random(tuple(base_shape)) - k = int(rng.integers(z.size)) - y = self._partition(z, k, xp, axis=None) - self._assert_valid_partition(z, k, y, xp, axis=None) - - def _test_input_validation(self, xp: ArrayNamespace): - with pytest.raises(TypeError): - _ = self._partition(np.asarray(1), 1, xp) - with pytest.raises(ValueError, match="out of bounds"): - _ = self._partition(np.asarray([1, 2]), 3, xp) - - def test_1d(self, xp: ArrayNamespace): - self._test_1d(xp) - - @pytest.mark.parametrize("ndim", [2, 3, 4]) - def test_nd(self, xp: ArrayNamespace, ndim: int): - self._test_nd(xp, ndim) - - def test_input_validation(self, xp: ArrayNamespace): - self._test_input_validation(xp) - - -@pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no argsort") -class TestArgpartition(TestPartition): - @classmethod - @override - def _partition( - cls, x: np.ndarray, k: int, xp: ArrayNamespace, axis: int | None = -1 - ): - arr = xp.asarray(x) - indices = argpartition(arr, k, axis=axis) - if axis is None: - arr = xp.reshape(arr, shape=(-1,)) - return arr[indices] - if arr.ndim == 1: - return arr[indices] - return cls._take_along_axis(arr, indices, axis=axis, xp=xp) - - @classmethod - def _take_along_axis( - cls, arr: Array, indices: Array, axis: int, xp: ArrayNamespace - ): - if hasattr(xp, "take_along_axis"): - return xp.take_along_axis(arr, indices, axis=axis) - if arr.ndim == 1: - return arr[indices] - if axis == 0: - assert isinstance(arr.shape[1], int) - arrs = [] - for i in range(arr.shape[1]): - arrs.append( - cls._take_along_axis( - arr[:, i, ...], indices[:, i, ...], axis=0, xp=xp - ) - ) - return xp.stack(arrs, axis=1) - axis = axis - 1 if axis != -1 else -1 - assert isinstance(arr.shape[0], int) - arrs = [] - for i in range(arr.shape[0]): - arrs.append( - cls._take_along_axis(arr[i, ...], indices[i, ...], axis=axis, xp=xp) - ) - return xp.stack(arrs, axis=0) - - @override - def test_1d(self, xp: ArrayNamespace): - self._test_1d(xp) - - @pytest.mark.parametrize("ndim", [2, 3, 4]) - @override - def test_nd(self, xp: ArrayNamespace, ndim: int): - self._test_nd(xp, ndim) - - @override - def test_input_validation(self, xp: ArrayNamespace): - self._test_input_validation(xp) - - -@pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no unique_inverse") -class TestIsIn: - def test_simple(self, xp: ArrayNamespace, library: Backend): - if library.like(Backend.NUMPY) and NUMPY_VERSION < (1, 24): - pytest.xfail("NumPy <1.24 has no kind kwarg in isin") - - b = xp.asarray([1, 2, 3, 4]) - - # `a` with 1 dimension - a = xp.asarray([1, 3, 6, 10]) - expected = xp.asarray([True, True, False, False]) - res = isin(a, b) - assert_equal(res, expected) - - # `a` with 2 dimensions - a = xp.asarray([[0, 2], [4, 6]]) - expected = xp.asarray([[False, True], [True, False]]) - res = isin(a, b) - assert_equal(res, expected) - - def test_device(self, xp: ArrayNamespace, device: Device, library: Backend): - if library.like(Backend.NUMPY) and NUMPY_VERSION < (1, 24): - pytest.xfail("NumPy <1.24 has no kind kwarg in isin") - - a = xp.asarray([1, 3, 6], device=device) - b = xp.asarray([1, 2, 3], device=device) - assert get_device(isin(a, b)) == device - - def test_assume_unique_and_invert( - self, xp: ArrayNamespace, device: Device, library: Backend - ): - if library.like(Backend.NUMPY) and NUMPY_VERSION < (1, 24): - pytest.xfail("NumPy <1.24 has no kind kwarg in isin") - - a = xp.asarray([0, 3, 6, 10], device=device) - b = xp.asarray([1, 2, 3, 10], device=device) - expected = xp.asarray([True, False, True, False], device=device) - res = isin(a, b, assume_unique=True, invert=True) - assert get_device(res) == device - assert_equal(res, expected) - - def test_kind(self, xp: ArrayNamespace, library: Backend): - if library.like(Backend.NUMPY) and NUMPY_VERSION < (1, 24): - pytest.xfail("NumPy <1.24 has no kind kwarg in isin") - - a = xp.asarray([0, 3, 6, 10]) - b = xp.asarray([1, 2, 3, 10]) - expected = xp.asarray([False, True, False, True]) - res = isin(a, b, kind="sort") - assert_equal(res, expected) - - -def _apply_over_batch(*argdefs: tuple[str, int]) -> Any: - """ - Factory for decorator that applies a function over batched arguments. - - Copied (with light simplifications) from `scipy._lib._util`. - - Array arguments may have any number of core dimensions (typically 0, - 1, or 2) and any broadcastable batch shapes. There may be any - number of array outputs of any number of dimensions. Assumptions - right now - which are satisfied by all functions of interest in `linalg` - - are that all array inputs are consecutive keyword or positional arguments, - and that the wrapped function returns either a single array or a tuple of - arrays. It's only as general as it needs to be right now - it can be extended. - - Parameters - ---------- - *argdefs : tuple of (str, int) - Definitions of array arguments: the keyword name of the argument, and - the number of core dimensions. - - Example: - -------- - `linalg.eig` accepts two matrices as the first two arguments `a` and `b`, where - `b` is optional, and returns one array or a tuple of arrays, depending on the - values of other positional or keyword arguments. To generate a wrapper that applies - the function over batches of `a` and optionally `b` : - - >>> _apply_over_batch(('a', 2), ('b', 2)) - """ - names, ndims = list(zip(*argdefs, strict=True)) - n_arrays = len(names) - - def decorator(f: Any) -> Any: - def wrapper( - *args_tuple: Any, - **kwargs: Any, - ) -> Any: - args = list(args_tuple) - - # Ensure all arrays in `arrays`, other arguments in `other_args`/`kwargs` - arrays, other_args = args[:n_arrays], args[n_arrays:] - arrays = cast(list[Array | None], arrays) - for i, name in enumerate(names): - if name in kwargs: - if i + 1 <= len(args): - message = ( - f"{f.__name__}() got multiple values for argument `{name}`." - ) - raise ValueError(message) - arrays.append(kwargs.pop(name)) - - xp = array_namespace(*arrays) - - # Determine core and batch shapes - batch_shapes = [] - core_shapes = [] - for i, (array, ndim) in enumerate(zip(arrays, ndims, strict=True)): - array = None if array is None else xp.asarray(array) # noqa: PLW2901 - shape = () if array is None else array.shape - arrays[i] = array - batch_shapes.append(shape[:-ndim] if ndim > 0 else shape) - core_shapes.append(shape[-ndim:] if ndim > 0 else ()) - - # Early exit if call is not batched - if not any(batch_shapes): - return f(*arrays, *other_args, **kwargs) - - # Determine broadcasted batch shape - batch_shape = np.broadcast_shapes(*batch_shapes) # Gives OK error message - - # Broadcast arrays to appropriate shape - for i, (array, core_shape) in enumerate( - zip(arrays, core_shapes, strict=True) - ): - if array is None: - continue - arrays[i] = xp.broadcast_to(array, batch_shape + core_shape) - - # Main loop - results = [] - for index in np.ndindex(batch_shape): - result = f( - *( - (array[index] if array is not None else None) - for array in arrays - ), - *other_args, - **kwargs, - ) - # Assume `result` is either a tuple or single array. This is easily - # generalized by allowing the contributor to pass an `unpack_result` - # callable to the decorator factory. - result = (result,) if not isinstance(result, tuple) else result - results.append(result) - results = list(zip(*results, strict=True)) - - # Reshape results - for i, result in enumerate(results): - result = xp.stack(result) # noqa: PLW2901 - core_shape = result.shape[1:] - results[i] = xp.reshape(result, batch_shape + core_shape) - - # Assume `result` should be a single array if there is only one element or - # a `tuple` otherwise. This is easily generalized by allowing the - # contributor to pass an `pack_result` callable to the decorator factory. - return results[0] if len(results) == 1 else results - - return wrapper - - return decorator - - -@_apply_over_batch(("a", 1), ("v", 1)) # type: ignore[untyped-decorator] -def xp_searchsorted( - a: Array, - v: Array, - side: Literal["left", "right"], - xp: ArrayNamespace, -) -> Array: - return xp.searchsorted(a, v, side=side) - - -@pytest.mark.skip_xp_backend(Backend.DASK, reason="no take_along_axis") -@pytest.mark.skip_xp_backend(Backend.SPARSE, reason="no searchsorted") -class TestSearchsorted: - def test_input_validation(self, xp: ArrayNamespace): - message = "`side` must be either 'left' or 'right'." - with pytest.raises(ValueError, match=message): - _ = xpx_searchsorted(xp.asarray([1, 2]), xp.asarray([1, 2]), side="center") # type: ignore[arg-type] # pyright: ignore[reportArgumentType] - - @pytest.mark.parametrize("side", ["left", "right"]) - @pytest.mark.parametrize("ties", [False, True]) - @pytest.mark.parametrize( - "shape", [0, 1, 2, 10, 11, 1000, 10001, (2, 0), (0, 2), (2, 10), (2, 3, 11)] - ) - @pytest.mark.parametrize("nans_x", [False, True]) - @pytest.mark.parametrize("infs_x", [False, True]) - @pytest.mark.parametrize("searchsorted", [xpx_searchsorted, _funcs_searchsorted]) - def test_nd( - self, - side: Literal["left", "right"], - ties: bool, - shape: int | tuple[int], - nans_x: bool, - infs_x: bool, - xp: ArrayNamespace, - searchsorted: Callable[..., Array], - ): - if nans_x and is_jax_namespace(xp): - pytest.xfail("https://github.com/jax-ml/jax/issues/39887") - if nans_x and is_torch_namespace(xp) and searchsorted == xpx_searchsorted: - pytest.skip("torch sorts NaNs differently") - if isinstance(shape, tuple) and searchsorted == _funcs_searchsorted: - message = ( - "Redundant; `xpx_searchsorted` delegates to " - "`_funcs_searchsorted` for multidimensional input." - ) - pytest.skip(message) - rng = np.random.default_rng(945298725498274853) - x = rng.integers(5, size=shape) if ties else rng.random(shape) - # float32 is to accommodate JAX - nextafter with `float64` is too small? - x = np.asarray(x, dtype=np.float32) # type:ignore[assignment] - xr = np.nextafter(x, np.inf) - xl = np.nextafter(x, -np.inf) - x_ = np.asarray([-np.inf, np.inf, np.nan]) - x_ = np.broadcast_to(x_, (*x.shape[:-1], 3)) - y = rng.permuted(np.concatenate((xl, x, xr, x_), axis=-1), axis=-1) - if nans_x: - mask = rng.random(shape) < 0.1 - x[mask] = np.nan - if infs_x: - mask = rng.random(shape) < 0.1 - x[mask] = -np.inf - mask = rng.random(shape) > 0.9 - x[mask] = np.inf - x = np.sort(x, axis=-1) # type:ignore[assignment] - x, y = np.asarray(x, dtype=np.float64), np.asarray(y, dtype=np.float64) - xp_default_int = default_dtype(xp, kind="integral") - if x.size == 0 and x.ndim > 0 and x.shape[-1] != 0: - ref = xp.empty((*x.shape[:-1], y.shape[-1]), dtype=xp_default_int) - else: - ref = xp_searchsorted(x, y, side=side, xp=np) - ref = xp.asarray(ref, dtype=xp_default_int) - x, y = xp.asarray(x.copy()), xp.asarray(y.copy()) - res = searchsorted(x, y, side=side, xp=xp) - assert_equal(res, ref) - - -@pytest.mark.skip_xp_backend( - Backend.ARRAY_API_STRICTEST, - reason="data_dependent_shapes flag for unique_values is disabled", -) -class TestUnion1d: - def test_simple(self, xp: ArrayNamespace): - a = xp.asarray([-1, 1, 0]) - b = xp.asarray([2, -2, 0]) - expected = xp.asarray([-2, -1, 0, 1, 2]) - res = union1d(a, b) - assert_equal(res, expected) - - def test_2d(self, xp: ArrayNamespace): - a = xp.asarray([[-1, 1, 0], [1, 2, 0]]) - b = xp.asarray([[1, 0, 1], [-2, -1, 0]]) - expected = xp.asarray([-2, -1, 0, 1, 2]) - res = union1d(a, b) - assert_equal(res, expected) - - def test_3d(self, xp: ArrayNamespace): - a = xp.asarray([[[-1, 0], [1, 2]], [[-1, 0], [1, 2]]]) - b = xp.asarray([[[0, 1], [-1, 2]], [[1, -2], [0, 2]]]) - expected = xp.asarray([-2, -1, 0, 1, 2]) - res = union1d(a, b) - assert_equal(res, expected) - - @pytest.mark.skip_xp_backend(Backend.TORCH, reason="materialize 'meta' device") - def test_device(self, xp: ArrayNamespace, device: Device): - a = xp.asarray([-1, 1, 0], device=device) - b = xp.asarray([2, -2, 0], device=device) - assert get_device(union1d(a, b)) == device - - -class TestAngle: - def test_simple(self, xp: ArrayNamespace): - a = xp.asarray([1, 0]) - res = angle(a) - expected = xp.asarray([0.0, 0.0], dtype=res.dtype) - assert_equal(res, expected) - - def test_basic(self, xp: ArrayNamespace): - x = xp.asarray( - [ - 1 + 3j, - np.sqrt(2) / 2.0 + 1j * np.sqrt(2) / 2, - 1, - 1j, - -1, - -1j, - 1 - 3j, - -1 + 3j, - ], - dtype=xp.complex128, - ) - expected = xp.asarray( - [ - np.arctan(3.0 / 1.0), - np.arctan(1.0), - 0, - np.pi / 2, - np.pi, - -np.pi / 2.0, - -np.arctan(3.0 / 1.0), - np.pi - np.arctan(3.0 / 1.0), - ], - dtype=xp.float64, - ) - assert_close(angle(x, xp=xp), expected, rtol=0, atol=1e-11) - assert_close( - angle(x, deg=True, xp=xp), - expected * 180 / xp.pi, - rtol=0, - atol=1e-11, - ) - - def test_real(self, xp: ArrayNamespace): - x = xp.asarray([0.0, -0.0, 1.0, -1.0]) - expected = xp.asarray([0.0, xp.pi, 0.0, xp.pi], dtype=x.dtype) - assert_close(angle(x, xp=xp), expected) - - def test_complex(self, xp: ArrayNamespace): - a = xp.asarray([1 + 1j, 1 - 1j, -1 + 1j, -1 - 1j]) - expected = xp.asarray([xp.pi / 4, -xp.pi / 4, 3 * xp.pi / 4, -3 * xp.pi / 4]) - res = angle(a, xp=xp) - assert_equal(res, expected) - - def test_integral(self, xp: ArrayNamespace): - x = xp.asarray([0, -1, 1], dtype=xp.int32) - actual = angle(x, xp=xp) - expected = xp.asarray( - [0.0, xp.pi, 0.0], dtype=default_dtype(xp, device=get_device(x)) - ) - assert_close(actual, expected) - - def test_2d(self, xp: ArrayNamespace): - a = xp.asarray([[1 + 1j, 1 - 1j], [-1 + 1j, -1 - 1j]]) - expected = xp.asarray( - [[xp.pi / 4, -xp.pi / 4], [3 * xp.pi / 4, -3 * xp.pi / 4]] - ) - res = angle(a, xp=xp) - assert_equal(res, expected) - - @pytest.mark.skip_xp_backend(Backend.TORCH, reason="materialize 'meta' device") - def test_device(self, xp: ArrayNamespace, device: Device): - a = xp.asarray([1 + 1j], device=device) - assert get_device(angle(a)) == device - - -class TestDeg2Rad: - def test_basic(self, xp: ArrayNamespace): - x = xp.asarray([0.0, 90.0, 180.0, 270.0, 360.0]) - expected = xp.asarray([0.0, xp.pi / 2, xp.pi, 3 * xp.pi / 2, 2 * xp.pi]) - assert_close(deg2rad(x), expected) - - @pytest.mark.parametrize("dtype_name", ["int32", "int64"]) - def test_integral(self, xp: ArrayNamespace, dtype_name: str): - x = xp.asarray([0, 90, 180], dtype=getattr(xp, dtype_name)) - actual = deg2rad(x, xp=xp) - expected = xp.asarray( - [0.0, xp.pi / 2, xp.pi], dtype=default_dtype(xp, device=get_device(x)) - ) - assert actual.dtype == expected.dtype - assert_close(actual, expected) - - def test_complex(self, xp: ArrayNamespace): - x = xp.asarray([180 + 90j], dtype=xp.complex64) - actual = deg2rad(x, xp=xp) - expected = xp.asarray([xp.pi + xp.pi / 2 * 1j], dtype=x.dtype) - assert actual.dtype == x.dtype - assert_close(actual, expected) - - def test_bool(self, xp: ArrayNamespace): - x = xp.asarray([True]) - with pytest.raises(TypeError, match="integral, real floating, or complex"): - _ = deg2rad(x, xp=xp) - - def test_device(self, xp: ArrayNamespace, device: Device): - x = xp.asarray([0.0, 90.0, 180.0], device=device) - assert get_device(deg2rad(x)) == device - - -class TestRad2Deg: - def test_basic(self, xp: ArrayNamespace): - x = xp.asarray([0.0, xp.pi / 2, xp.pi, 3 * xp.pi / 2, 2 * xp.pi]) - expected = xp.asarray([0.0, 90.0, 180.0, 270.0, 360.0]) - assert_close(rad2deg(x), expected) - - @pytest.mark.parametrize("dtype_name", ["int32", "int64"]) - def test_integral(self, xp: ArrayNamespace, dtype_name: str): - x = xp.asarray([0, 1, 2], dtype=getattr(xp, dtype_name)) - actual = rad2deg(x, xp=xp) - expected = xp.asarray( - [0.0, 180 / xp.pi, 360 / xp.pi], - dtype=default_dtype(xp, device=get_device(x)), - ) - assert actual.dtype == expected.dtype - assert_close(actual, expected) - - def test_complex(self, xp: ArrayNamespace): - x = xp.asarray([xp.pi + xp.pi / 2 * 1j], dtype=xp.complex64) - actual = rad2deg(x, xp=xp) - expected = xp.asarray([180 + 90j], dtype=x.dtype) - assert actual.dtype == x.dtype - assert_close(actual, expected) - - def test_bool(self, xp: ArrayNamespace): - x = xp.asarray([True]) - with pytest.raises(TypeError, match="integral, real floating, or complex"): - _ = rad2deg(x, xp=xp) - - def test_device(self, xp: ArrayNamespace, device: Device): - x = xp.asarray([0.0, xp.pi / 2, xp.pi], device=device) - assert get_device(rad2deg(x)) == device - - -class TestUnravelIndex: - def test_simple(self, xp: ArrayNamespace): - indices = xp.asarray([22, 41, 37]) - shape = (7, 6) - expected = (xp.asarray([3, 6, 6]), xp.asarray([4, 5, 1])) - res = unravel_index(indices, shape) - for res_arr, exp_arr in zip(res, expected, strict=True): - assert_equal(res_arr, exp_arr) - - indices = xp.asarray([0, 1, 2, 3, 4, 5]) - shape = (3, 2) - expected = ( - xp.asarray([0, 0, 1, 1, 2, 2]), - xp.asarray([0, 1, 0, 1, 0, 1]), - ) - res = unravel_index(indices, shape) - for res_arr, exp_arr in zip(res, expected, strict=True): - assert_equal(res_arr, exp_arr) - - def test_indices_scalar(self, xp: ArrayNamespace): - indices = xp.asarray(1621) - shape = (6, 7, 8, 9) - expected = (xp.asarray(3), xp.asarray(1), xp.asarray(4), xp.asarray(1)) - res = unravel_index(indices, shape) - # a tuple of integers is expected - assert res == expected - - def test_indices_2d(self, xp: ArrayNamespace): - indices = xp.asarray([[1234], [5678]]) - shape = (10, 10, 10, 10) - expected = ( - xp.asarray([[1], [5]]), - xp.asarray([[2], [6]]), - xp.asarray([[3], [7]]), - xp.asarray([[4], [8]]), - ) - res = unravel_index(indices, shape) - for res_arr, exp_arr in zip(res, expected, strict=True): - assert_equal(res_arr, exp_arr) - - def test_device(self, xp: ArrayNamespace, device: Device): - indices = xp.asarray([4, 1], device=device) - shape = (3, 2) - res = unravel_index(indices, shape) - for res_arr in res: - assert get_device(res_arr) == device - - @pytest.mark.skip_xp_backend(Backend.NUMPY_READONLY, reason="xp=xp") - def test_xp(self, xp: ArrayNamespace): - indices = xp.asarray([1, 5]) - shape = (3, 2) - expected = ( - xp.asarray([0, 2]), - xp.asarray([1, 1]), - ) - res = unravel_index(indices, shape, xp=xp) - for res_arr, exp_arr in zip(res, expected, strict=True): - assert_equal(res_arr, exp_arr) - - -class TestNanMin: - def test_simple(self, xp: ArrayNamespace): - a = xp.asarray([[1, 2], [3, xp.nan]]) - - # with the default `axis=None` a single scalar is returned - res = nanmin(a) - expected = 1.0 - assert res == expected - - res = nanmin(a, axis=0) - expected = xp.asarray([1.0, 2.0]) - assert_equal(res, expected) - - res = nanmin(a, axis=1) - expected = xp.asarray([1.0, 3.0]) - assert_equal(res, expected) - - def test_bigger(self, xp: ArrayNamespace): - a = xp.asarray( - [ - [1, xp.nan, 4, 5], - [xp.nan, -2, xp.nan, -4], - [2, 1, 3, xp.nan], - ] - ) - - res = nanmin(a, axis=0) - expected = xp.asarray([1.0, -2.0, 3.0, -4.0]) - assert_equal(res, expected) - - res = nanmin(a, axis=1) - expected = xp.asarray([1.0, -4.0, 1.0]) - assert_equal(res, expected) - - def test_with_infinity(self, xp: ArrayNamespace): - a = xp.asarray([0.1, 1.0, xp.nan, xp.inf]) - res = nanmin(a) - expected = 0.1 - assert res == expected - - a = xp.asarray([0.1, 1.0, xp.nan, -xp.inf]) - res = nanmin(a) - expected = -xp.inf - assert res == expected - - def test_scalar(self, xp: ArrayNamespace): - a = xp.asarray(1.0) - assert nanmin(a) == 1.0 - - @pytest.mark.filterwarnings("ignore:.*All-NaN slice*.:RuntimeWarning") - def test_all_nan_slice_2d(self, xp: ArrayNamespace): - a = xp.asarray( - [ - [xp.nan, 5.0], - [xp.nan, 2.0], - ] - ) - - res = nanmin(a, axis=0, xp=xp) - expected = xp.asarray([xp.nan, 2.0]) - assert_equal(res, expected) - - @pytest.mark.skip_xp_backend( - Backend.TORCH, reason="torch.nanmin does not support tensors on meta device" - ) - @pytest.mark.parametrize("axis", [None, 0, 1]) - def test_device(self, axis: int | None, xp: ArrayNamespace, device: Device): - a = xp.asarray([[4, xp.nan, 1], [2, 5, xp.nan]], device=device) - res = nanmin(a, axis=axis) - assert get_device(res) == device - - @pytest.mark.parametrize( - ("axis", "expected_list"), [(0, [2.0, 3.0, 1.0]), (1, [1.0, 2.0])] - ) - def test_xp(self, axis: int | None, expected_list: list[float], xp: ArrayNamespace): - a = xp.asarray([[4, xp.nan, 1], [2, 3, xp.nan]]) - res = nanmin(a, axis=axis, xp=xp) - expected = xp.asarray(expected_list) - assert_equal(res, expected) - - -class TestNanMax: - def test_simple(self, xp: ArrayNamespace): - a = xp.asarray([[5, 3], [6, xp.nan]]) - - # with the default `axis=None` a single scalar is returned - res = nanmax(a) - expected = 6.0 - assert res == expected - - res = nanmax(a, axis=0) - expected = xp.asarray([6.0, 3.0]) - assert_equal(res, expected) - - res = nanmax(a, axis=1) - expected = xp.asarray([5.0, 6.0]) - assert_equal(res, expected) - - def test_bigger(self, xp: ArrayNamespace): - a = xp.asarray( - [ - [1, xp.nan, 4, 5], - [xp.nan, 2, xp.nan, 4], - [6, 1, 3, xp.nan], - ] - ) - - res = nanmax(a, axis=0) - expected = xp.asarray([6.0, 2.0, 4.0, 5.0]) - assert_equal(res, expected) - - res = nanmax(a, axis=1) - expected = xp.asarray([5.0, 4.0, 6.0]) - assert_equal(res, expected) - - def test_with_infinity(self, xp: ArrayNamespace): - a = xp.asarray([0.1, 5.0, xp.nan, -xp.inf]) - res = nanmax(a) - expected = 5.0 - assert res == expected - - a = xp.asarray([3.0, 10.0, xp.nan, xp.inf]) - res = nanmax(a) - expected = xp.inf - assert res == expected - - def test_scalar(self, xp: ArrayNamespace): - a = xp.asarray(1.0) - assert nanmax(a) == 1.0 - - @pytest.mark.filterwarnings("ignore:.*All-NaN slice*.:RuntimeWarning") - def test_all_nan_slice_2d(self, xp: ArrayNamespace): - a = xp.asarray( - [ - [xp.nan, 5.0], - [xp.nan, 2.0], - ] - ) - - res = nanmax(a, axis=0, xp=xp) - expected = xp.asarray([xp.nan, 5.0]) - assert_equal(res, expected) - - @pytest.mark.skip_xp_backend( - Backend.TORCH, reason="torch.nanmax does not support tensors on meta device" - ) - @pytest.mark.parametrize("axis", [None, 0, 1]) - def test_device(self, axis: int | None, xp: ArrayNamespace, device: Device): - a = xp.asarray([[4, xp.nan, 1], [2, 5, xp.nan]], device=device) - res = nanmax(a, axis=axis) - assert get_device(res) == device - - @pytest.mark.parametrize( - ("axis", "expected_list"), [(0, [4.0, 3.0, 1.0]), (1, [4.0, 3.0])] - ) - def test_xp(self, axis: int | None, expected_list: list[float], xp: ArrayNamespace): - a = xp.asarray([[4, xp.nan, 1], [2, 3, xp.nan]]) - res = nanmax(a, axis=axis, xp=xp) - expected = xp.asarray(expected_list) - assert_equal(res, expected) - - -class TestNanSum: - def test_simple(self, xp: ArrayNamespace): - a = xp.asarray([[1.0, 2.0], [3.0, xp.nan]]) - - res = nansum(a) - expected = 6.0 - assert res == expected - - res = nansum(a, axis=0) - expected = xp.asarray([4.0, 2.0]) - assert_equal(res, expected) - - res = nansum(a, axis=1) - expected = xp.asarray([3.0, 3.0]) - assert_equal(res, expected) - - def test_bigger(self, xp: ArrayNamespace): - a = xp.asarray( - [ - [1.0, xp.nan, 4.0, 5.0], - [xp.nan, -2.0, xp.nan, -4.0], - [2.0, 1.0, 3.0, xp.nan], - ] - ) - - res = nansum(a, axis=0) - expected = xp.asarray([3.0, -1.0, 7.0, 1.0]) - assert_equal(res, expected) - - res = nansum(a, axis=1) - expected = xp.asarray([10.0, -6.0, 6.0]) - assert_equal(res, expected) - - def test_all_nan_slice(self, xp: ArrayNamespace): - a = xp.asarray([[xp.nan, 1.0], [xp.nan, xp.nan]]) - - res = nansum(a, axis=0) - expected = xp.asarray([0.0, 1.0]) - assert_equal(res, expected) - - def test_scalar(self, xp: ArrayNamespace): - a = xp.asarray(1.0) - assert nansum(a) == 1.0 - - @pytest.mark.skip_xp_backend( - Backend.TORCH, reason="torch.nansum does not support tensors on meta device" - ) - @pytest.mark.parametrize("axis", [None, 0, 1]) - def test_device(self, axis: int | None, xp: ArrayNamespace, device: Device): - a = xp.asarray([[4.0, xp.nan, 1.0], [2.0, 5.0, xp.nan]], device=device) - res = nansum(a, axis=axis) - assert get_device(res) == device - - @pytest.mark.parametrize( - ("axis", "expected_list"), [(0, [6.0, 3.0, 1.0]), (1, [5.0, 5.0])] - ) - def test_xp(self, axis: int | None, expected_list: list[float], xp: ArrayNamespace): - a = xp.asarray([[4.0, xp.nan, 1.0], [2.0, 3.0, xp.nan]]) - res = nansum(a, axis=axis, xp=xp) - expected = xp.asarray(expected_list) - assert_equal(res, expected) diff --git a/tests/test_helpers.py b/tests/test_helpers.py index 215f0593..c8a93de0 100644 --- a/tests/test_helpers.py +++ b/tests/test_helpers.py @@ -5,9 +5,9 @@ import pytest from array_api_extra._lib._backends import Backend -from array_api_extra._lib._utils._compat import array_namespace -from array_api_extra._lib._utils._compat import device as get_device -from array_api_extra._lib._utils._helpers import ( +from array_api_extra._lib._compat import array_namespace +from array_api_extra._lib._compat import device as get_device +from array_api_extra._lib._helpers import ( asarrays, capabilities, eager_shape, @@ -18,7 +18,7 @@ pickle_flatten, pickle_unflatten, ) -from array_api_extra._lib._utils._typing import Array, ArrayNamespace, Device, DType +from array_api_extra._lib._typing import Array, ArrayNamespace, Device, DType from array_api_extra.testing import assert_equal, lazy_xp_function from .conftest import np_compat diff --git a/tests/test_indexing.py b/tests/test_indexing.py new file mode 100644 index 00000000..4b2c83ae --- /dev/null +++ b/tests/test_indexing.py @@ -0,0 +1,201 @@ +from collections.abc import Callable + +import numpy as np +import pytest + +from array_api_extra import diag_indices, tril_indices, triu_indices, unravel_index +from array_api_extra._lib._backends import Backend +from array_api_extra._lib._compat import device as get_device +from array_api_extra._lib._typing import Array, ArrayNamespace, Device +from array_api_extra.testing import assert_equal, lazy_xp_function + +lazy_xp_function(diag_indices) +lazy_xp_function(tril_indices) +lazy_xp_function(triu_indices) + + +@pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no arange", strict=False) +class TestDiagIndices: + def test_basic(self, xp: ArrayNamespace): + rows, cols = diag_indices(5, xp=xp) + ref_rows, ref_cols = np.diag_indices(5) + assert_equal(rows, xp.asarray(ref_rows)) + assert_equal(cols, xp.asarray(ref_cols)) + + @pytest.mark.parametrize("n", [2, 4, 7]) + @pytest.mark.parametrize("ndim", [1, 2, 3, 4]) + def test_ndim(self, xp: ArrayNamespace, n: int, ndim: int): + idx = diag_indices(n, ndim=ndim, xp=xp) + assert len(idx) == ndim + ref = np.diag_indices(n, ndim=ndim) + for got, expected in zip(idx, ref, strict=True): + assert_equal(got, xp.asarray(expected)) + + def test_empty(self, xp: ArrayNamespace): + rows, cols = diag_indices(0, xp=xp) + assert rows.shape == (0,) + assert cols.shape == (0,) + + def test_validation(self, xp: ArrayNamespace): + with pytest.raises(ValueError, match="`n` must be non-negative"): + _ = diag_indices(-1, xp=xp) + with pytest.raises(ValueError, match="`ndim` must be >= 1"): + _ = diag_indices(3, ndim=0, xp=xp) + + def test_device(self, xp: ArrayNamespace, device: Device): + default_device = get_device(xp.empty(0)) + rows, cols = diag_indices(3, device=None, xp=xp) + assert get_device(rows) == default_device + assert get_device(cols) == default_device + rows, cols = diag_indices(3, device=device, xp=xp) + assert get_device(rows) == device + assert get_device(cols) == device + + +@pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no arange/nonzero", strict=False) +@pytest.mark.xfail_xp_backend( + Backend.ARRAY_API_STRICTEST, + reason="generic path uses nonzero (data-dependent)", + strict=False, +) +@pytest.mark.parametrize( + ("xpx_fn", "np_fn"), + [(tril_indices, np.tril_indices), (triu_indices, np.triu_indices)], + ids=["tril", "triu"], +) +class TestTriIndices: + def test_basic( + self, + xp: ArrayNamespace, + xpx_fn: Callable[..., tuple[Array, Array]], + np_fn: Callable[..., tuple[Array, Array]], + ): + rows, cols = xpx_fn(4, xp=xp) + ref_rows, ref_cols = np_fn(4) + assert_equal(rows, xp.asarray(ref_rows)) + assert_equal(cols, xp.asarray(ref_cols)) + + @pytest.mark.parametrize("offset", [-2, -1, 0, 1, 2]) + def test_offset( + self, + xp: ArrayNamespace, + xpx_fn: Callable[..., tuple[Array, Array]], + np_fn: Callable[..., tuple[Array, Array]], + offset: int, + ): + rows, cols = xpx_fn(5, offset=offset, xp=xp) + ref_rows, ref_cols = np_fn(5, k=offset) + assert_equal(rows, xp.asarray(ref_rows)) + assert_equal(cols, xp.asarray(ref_cols)) + + def test_rectangular( + self, + xp: ArrayNamespace, + xpx_fn: Callable[..., tuple[Array, Array]], + np_fn: Callable[..., tuple[Array, Array]], + ): + rows, cols = xpx_fn(3, m=5, xp=xp) + ref_rows, ref_cols = np_fn(3, m=5) + assert_equal(rows, xp.asarray(ref_rows)) + assert_equal(cols, xp.asarray(ref_cols)) + + @pytest.mark.xfail_xp_backend( + Backend.DASK, reason="dask: no 2D fancy indexing", strict=False + ) + def test_use_to_read( + self, + xp: ArrayNamespace, + xpx_fn: Callable[..., tuple[Array, Array]], + np_fn: Callable[..., tuple[Array, Array]], + ): + rng = np.random.default_rng(0) + a = rng.integers(0, 100, (4, 4)) + a_xp = xp.asarray(a) + rows, cols = xpx_fn(4, xp=xp) + assert_equal(a_xp[rows, cols], xp.asarray(a[np_fn(4)])) + + def test_validation( + self, + xp: ArrayNamespace, + xpx_fn: Callable[..., tuple[Array, Array]], + np_fn: Callable[..., tuple[Array, Array]], # noqa: ARG002 # pytest param + ): + with pytest.raises(ValueError, match="`n` must be non-negative"): + _ = xpx_fn(-1, xp=xp) + with pytest.raises(ValueError, match="`m` must be non-negative"): + _ = xpx_fn(3, m=-1, xp=xp) + + def test_device( + self, + xp: ArrayNamespace, + device: Device, + xpx_fn: Callable[..., tuple[Array, Array]], + np_fn: Callable[..., tuple[Array, Array]], # noqa: ARG002 # pytest param + ): + default_device = get_device(xp.empty(0)) + rows, cols = xpx_fn(4, device=None, xp=xp) + assert get_device(rows) == default_device + assert get_device(cols) == default_device + rows, cols = xpx_fn(4, device=device, xp=xp) + assert get_device(rows) == device + assert get_device(cols) == device + + +class TestUnravelIndex: + def test_simple(self, xp: ArrayNamespace): + indices = xp.asarray([22, 41, 37]) + shape = (7, 6) + expected = (xp.asarray([3, 6, 6]), xp.asarray([4, 5, 1])) + res = unravel_index(indices, shape) + for res_arr, exp_arr in zip(res, expected, strict=True): + assert_equal(res_arr, exp_arr) + + indices = xp.asarray([0, 1, 2, 3, 4, 5]) + shape = (3, 2) + expected = ( + xp.asarray([0, 0, 1, 1, 2, 2]), + xp.asarray([0, 1, 0, 1, 0, 1]), + ) + res = unravel_index(indices, shape) + for res_arr, exp_arr in zip(res, expected, strict=True): + assert_equal(res_arr, exp_arr) + + def test_indices_scalar(self, xp: ArrayNamespace): + indices = xp.asarray(1621) + shape = (6, 7, 8, 9) + expected = (xp.asarray(3), xp.asarray(1), xp.asarray(4), xp.asarray(1)) + res = unravel_index(indices, shape) + # a tuple of integers is expected + assert res == expected + + def test_indices_2d(self, xp: ArrayNamespace): + indices = xp.asarray([[1234], [5678]]) + shape = (10, 10, 10, 10) + expected = ( + xp.asarray([[1], [5]]), + xp.asarray([[2], [6]]), + xp.asarray([[3], [7]]), + xp.asarray([[4], [8]]), + ) + res = unravel_index(indices, shape) + for res_arr, exp_arr in zip(res, expected, strict=True): + assert_equal(res_arr, exp_arr) + + def test_device(self, xp: ArrayNamespace, device: Device): + indices = xp.asarray([4, 1], device=device) + shape = (3, 2) + res = unravel_index(indices, shape) + for res_arr in res: + assert get_device(res_arr) == device + + @pytest.mark.skip_xp_backend(Backend.NUMPY_READONLY, reason="xp=xp") + def test_xp(self, xp: ArrayNamespace): + indices = xp.asarray([1, 5]) + shape = (3, 2) + expected = ( + xp.asarray([0, 2]), + xp.asarray([1, 1]), + ) + res = unravel_index(indices, shape, xp=xp) + for res_arr, exp_arr in zip(res, expected, strict=True): + assert_equal(res_arr, exp_arr) diff --git a/tests/test_inspection.py b/tests/test_inspection.py new file mode 100644 index 00000000..203d6169 --- /dev/null +++ b/tests/test_inspection.py @@ -0,0 +1,39 @@ +import pytest + +from array_api_extra import default_dtype +from array_api_extra._lib._typing import ArrayNamespace, Device +from array_api_extra.testing import lazy_xp_function + +lazy_xp_function(default_dtype) + + +class TestDefaultDType: + def test_basic(self, xp: ArrayNamespace): + assert default_dtype(xp) == xp.empty(0).dtype + + def test_kind(self, xp: ArrayNamespace): + assert default_dtype(xp, "real floating") == xp.empty(0).dtype + assert default_dtype(xp, "complex floating") == (xp.empty(0) * 1j).dtype + assert default_dtype(xp, "integral") == xp.int64 + assert default_dtype(xp, "indexing") == xp.int64 + + with pytest.raises(ValueError, match="Unknown kind"): + _ = default_dtype(xp, "foo") # type: ignore[arg-type] # pyright: ignore[reportArgumentType] + + def test_device(self, xp: ArrayNamespace, device: Device): + # Note: at the moment there are no known namespaces with + # device-specific default dtypes. + assert default_dtype(xp, device=None) == xp.empty(0).dtype + assert default_dtype(xp, device=device) == xp.empty(0).dtype + + def test_torch(self, torch: ArrayNamespace): + xp = torch + xp.set_default_dtype(xp.float64) + assert default_dtype(xp) == xp.float64 + assert default_dtype(xp, "real floating") == xp.float64 + assert default_dtype(xp, "complex floating") == xp.complex128 + + xp.set_default_dtype(xp.float32) + assert default_dtype(xp) == xp.float32 + assert default_dtype(xp, "real floating") == xp.float32 + assert default_dtype(xp, "complex floating") == xp.complex64 diff --git a/tests/test_lazy.py b/tests/test_lazy.py index 774881c9..7222ec99 100644 --- a/tests/test_lazy.py +++ b/tests/test_lazy.py @@ -6,11 +6,11 @@ import array_api_extra as xpx # Let some tests bypass lazy_xp_function from array_api_extra import lazy_apply +from array_api_extra._lib import _compat from array_api_extra._lib._backends import Backend -from array_api_extra._lib._utils import _compat -from array_api_extra._lib._utils._compat import array_namespace, is_dask_array -from array_api_extra._lib._utils._helpers import eager_shape -from array_api_extra._lib._utils._typing import Array, ArrayNamespace, Device +from array_api_extra._lib._compat import array_namespace, is_dask_array +from array_api_extra._lib._helpers import eager_shape +from array_api_extra._lib._typing import Array, ArrayNamespace, Device from array_api_extra.testing import assert_equal, lazy_xp_function lazy_xp_function(lazy_apply) diff --git a/tests/test_linalg.py b/tests/test_linalg.py new file mode 100644 index 00000000..64288385 --- /dev/null +++ b/tests/test_linalg.py @@ -0,0 +1,90 @@ +import pytest + +from array_api_extra import kron +from array_api_extra._lib._compat import device as get_device +from array_api_extra._lib._typing import ArrayNamespace, Device +from array_api_extra.testing import assert_equal, lazy_xp_function + +lazy_xp_function(kron) + + +class TestKron: + def test_basic(self, xp: ArrayNamespace): + # Using 0-dimensional array + a = xp.asarray(1) + b = xp.asarray([[1, 2], [3, 4]]) + assert_equal(kron(a, b), b) + assert_equal(kron(b, a), b) + + # Using 1-dimensional array + a = xp.asarray([3]) + b = xp.asarray([[1, 2], [3, 4]]) + k = xp.asarray([[3, 6], [9, 12]]) + assert_equal(kron(a, b), k) + assert_equal(kron(b, a), k) + + # Using 3-dimensional array + a = xp.asarray([[[1]], [[2]]]) + b = xp.asarray([[1, 2], [3, 4]]) + k = xp.asarray([[[1, 2], [3, 4]], [[2, 4], [6, 8]]]) + assert_equal(kron(a, b), k) + assert_equal(kron(b, a), k) + + def test_kron_smoke(self, xp: ArrayNamespace): + a = xp.ones((3, 3)) + b = xp.ones((3, 3)) + k = xp.ones((9, 9)) + + assert_equal(kron(a, b), k) + + @pytest.mark.parametrize( + ("shape_a", "shape_b"), + [ + ((1, 1), (1, 1)), + ((1, 2, 3), (4, 5, 6)), + ((2, 2), (2, 2, 2)), + ((1, 0), (1, 1)), + ((2, 0, 2), (2, 2)), + ((2, 0, 0, 2), (2, 0, 2)), + ], + ) + def test_kron_shape( + self, xp: ArrayNamespace, shape_a: tuple[int, ...], shape_b: tuple[int, ...] + ): + a = xp.ones(shape_a) + b = xp.ones(shape_b) + normalised_shape_a = xp.asarray( + (1,) * max(0, len(shape_b) - len(shape_a)) + shape_a + ) + normalised_shape_b = xp.asarray( + (1,) * max(0, len(shape_a) - len(shape_b)) + shape_b + ) + expected_shape = tuple( + int(dim) for dim in xp.multiply(normalised_shape_a, normalised_shape_b) + ) + + k = kron(a, b) + assert k.shape == expected_shape + + def test_python_scalar(self, xp: ArrayNamespace): + a = 1 + # Test no dtype promotion to xp.asarray(a); use b.dtype + b = xp.asarray([[1, 2], [3, 4]], dtype=xp.int16) + assert_equal(kron(a, b), b) + assert_equal(kron(b, a), b) + assert_equal(kron(1, 1, xp=xp), xp.asarray(1)) + + def test_all_python_scalars(self): + with pytest.raises(TypeError, match=r"array_namespace requires .* array input"): + _ = kron(1, 1) + + def test_device(self, xp: ArrayNamespace, device: Device): + x1 = xp.asarray([1, 2, 3], device=device) + x2 = xp.asarray([4, 5], device=device) + assert get_device(kron(x1, x2)) == device + + def test_xp(self, xp: ArrayNamespace): + a = xp.ones((3, 3)) + b = xp.ones((3, 3)) + k = xp.ones((9, 9)) + assert_equal(kron(a, b, xp=xp), k) diff --git a/tests/test_manipulation.py b/tests/test_manipulation.py new file mode 100644 index 00000000..f6e06d04 --- /dev/null +++ b/tests/test_manipulation.py @@ -0,0 +1,373 @@ +import math + +import numpy as np +import pytest + +from array_api_extra import atleast_nd, broadcast_shapes, expand_dims, pad +from array_api_extra._lib._backends import Backend +from array_api_extra._lib._compat import device as get_device +from array_api_extra._lib._typing import ArrayNamespace, Device +from array_api_extra.testing import assert_equal, lazy_xp_function + +lazy_xp_function(atleast_nd) +lazy_xp_function(broadcast_shapes) +lazy_xp_function(expand_dims) +lazy_xp_function(pad) + + +class TestAtLeastND: + def test_0D(self, xp: ArrayNamespace): + x = xp.asarray(1.0) + + y = atleast_nd(x, ndim=0) + assert_equal(y, x) + + y = atleast_nd(x, ndim=1) + assert_equal(y, xp.ones((1,))) + + y = atleast_nd(x, ndim=5) + assert_equal(y, xp.ones((1, 1, 1, 1, 1))) + + @pytest.mark.parametrize( + ("input_shape", "ndim", "expected_shape"), + [ + ((1,), 0, (1,)), + ((5,), 1, (5,)), + ((2,), 2, (1, 2)), + ((3,), 3, (1, 1, 3)), + ((2,), 5, (1, 1, 1, 1, 2)), + ], + ) + def test_1D_shapes( + self, + input_shape: tuple[int], + ndim: int, + expected_shape: tuple[int], + xp: ArrayNamespace, + ): + n = math.prod(input_shape) + x = xp.asarray(np.arange(n).reshape(input_shape)) + y = atleast_nd(x, ndim=ndim) + + assert y.shape == expected_shape + assert xp.sum(y) == int(n * (n - 1) / 2) + + def test_1D_values(self, xp: ArrayNamespace): + x = xp.asarray([0, 1]) + + y = atleast_nd(x, ndim=0) + assert_equal(y, x) + + y = atleast_nd(x, ndim=1) + assert_equal(y, x) + + y = atleast_nd(x, ndim=2) + assert_equal(y, xp.asarray([[0, 1]])) + + y = atleast_nd(x, ndim=5) + assert_equal(y, xp.asarray([[[[[0, 1]]]]])) + + @pytest.mark.parametrize( + ("input_shape", "ndim", "expected_shape"), + [ + ((2, 1), 0, (2, 1)), + ((5, 2), 1, (5, 2)), + ((2, 1), 2, (2, 1)), + ((3, 1), 3, (1, 3, 1)), + ((2, 8), 5, (1, 1, 1, 2, 8)), + ], + ) + def test_2D_shapes( + self, + input_shape: tuple[int], + ndim: int, + expected_shape: tuple[int], + xp: ArrayNamespace, + ): + n = math.prod(input_shape) + x = xp.asarray(np.arange(n).reshape(input_shape)) + y = atleast_nd(x, ndim=ndim) + + assert y.shape == expected_shape + assert xp.sum(y) == int(n * (n - 1) / 2) + + def test_2D_values(self, xp: ArrayNamespace): + x = xp.asarray([[3.0], [4.0]]) + + y = atleast_nd(x, ndim=0) + assert_equal(y, x) + + y = atleast_nd(x, ndim=2) + assert_equal(y, x) + + y = atleast_nd(x, ndim=3) + assert_equal(y, xp.asarray([[[3.0], [4.0]]])) + + y = atleast_nd(x, ndim=5) + assert_equal(y, xp.asarray([[[[[3.0], [4.0]]]]])) + + @pytest.mark.parametrize( + ("input_shape", "ndim", "expected_shape"), + [ + ((2, 1, 1), 0, (2, 1, 1)), + ((1, 5, 2), 1, (1, 5, 2)), + ((2, 1, 1), 2, (2, 1, 1)), + ((1, 3, 1), 3, (1, 3, 1)), + ((2, 8, 1), 5, (1, 1, 2, 8, 1)), + ], + ) + def test_3D_shapes( + self, + input_shape: tuple[int], + ndim: int, + expected_shape: tuple[int], + xp: ArrayNamespace, + ): + n = math.prod(input_shape) + x = xp.asarray(np.arange(n).reshape(input_shape)) + y = atleast_nd(x, ndim=ndim) + + assert y.shape == expected_shape + assert xp.sum(y) == int(n * (n - 1) / 2) + + def test_3D_values(self, xp: ArrayNamespace): + x = xp.asarray([[[3.0], [2.0]]]) + + y = atleast_nd(x, ndim=0) + assert_equal(y, x) + + y = atleast_nd(x, ndim=2) + assert_equal(y, x) + + y = atleast_nd(x, ndim=3) + assert_equal(y, x) + + y = atleast_nd(x, ndim=5) + assert_equal(y, xp.asarray([[[[[3.0], [2.0]]]]])) + + @pytest.mark.parametrize( + ("input_shape", "ndim", "expected_shape"), + [ + ((2, 1, 1, 2, 1), 0, (2, 1, 1, 2, 1)), + ((1, 5, 2, 3, 2), 2, (1, 5, 2, 3, 2)), + ((2, 1, 1, 5, 2), 5, (2, 1, 1, 5, 2)), + ((1, 3, 1, 2, 1), 6, (1, 1, 3, 1, 2, 1)), + ((2, 8, 1, 9, 8), 9, (1, 1, 1, 1, 2, 8, 1, 9, 8)), + ], + ) + def test_5D_shapes( + self, + input_shape: tuple[int], + ndim: int, + expected_shape: tuple[int], + xp: ArrayNamespace, + ): + n = math.prod(input_shape) + x = xp.asarray(np.arange(n).reshape(input_shape)) + y = atleast_nd(x, ndim=ndim) + + assert y.shape == expected_shape + assert xp.sum(y) == int(n * (n - 1) / 2) + + def test_5D_values(self, xp: ArrayNamespace): + x = xp.asarray([[[[[3.0]], [[2.0]]]]]) + + y = atleast_nd(x, ndim=0) + assert_equal(y, x) + + y = atleast_nd(x, ndim=4) + assert_equal(y, x) + + y = atleast_nd(x, ndim=5) + assert_equal(y, x) + + y = atleast_nd(x, ndim=6) + assert_equal(y, xp.asarray([[[[[[3.0]], [[2.0]]]]]])) + + y = atleast_nd(x, ndim=9) + assert_equal(y, xp.asarray([[[[[[[[[3.0]], [[2.0]]]]]]]]])) + + +@pytest.mark.filterwarnings("ignore:.*removed in v1.0.0.*:DeprecationWarning") +class TestBroadcastShapes: + def test_delegates_known_integer_shapes(self, monkeypatch: pytest.MonkeyPatch): + calls = [] + + def mock_broadcast_shapes(*shapes: tuple[int, ...]) -> tuple[int, ...]: + calls.append(shapes) + return (99,) + + monkeypatch.setattr(np, "broadcast_shapes", mock_broadcast_shapes) + + assert broadcast_shapes((2,), (1,), xp=np) == (99,) + assert calls == [((2,), (1,))] + + def test_fallback_without_xp(self, monkeypatch: pytest.MonkeyPatch): + def mock_broadcast_shapes(*_shapes: tuple[int, ...]) -> tuple[int, ...]: + msg = "Native delegation should not be used without xp" + raise AssertionError(msg) + + monkeypatch.setattr(np, "broadcast_shapes", mock_broadcast_shapes) + + assert broadcast_shapes((2,), (1,)) == (2,) + + @pytest.mark.skip_xp_backend(Backend.NUMPY_READONLY, reason="xp=xp") + def test_xp(self, xp: ArrayNamespace): + assert broadcast_shapes((2, 3), (2, 1), xp=xp) == (2, 3) + + @pytest.mark.parametrize( + "args", + [ + (), + ((),), + ((), ()), + ((1,),), + ((1,), (1,)), + ((2,), (1,)), + ((3, 1, 4), (2, 1)), + ((1, 1, 4), (2, 1)), + ((1,), ()), + ((), (2,), ()), + ((0,),), + ((0,), (1,)), + ((2, 0), (1, 1)), + ((2, 0, 3), (2, 1, 1)), + ], + ) + def test_simple(self, args: tuple[tuple[int, ...], ...]): + expect = np.broadcast_shapes(*args) + actual = broadcast_shapes(*args) + assert actual == expect + + @pytest.mark.parametrize( + "args", + [ + ((2,), (3,)), + ((2, 3), (1, 2)), + ((2,), (0,)), + ((2, 0, 2), (1, 3, 1)), + ], + ) + def test_fail(self, args: tuple[tuple[int, ...], ...]): + match = "cannot be broadcast to a single shape" + with pytest.raises(ValueError, match=match): + _ = np.broadcast_shapes(*args) + with pytest.raises(ValueError, match=match): + _ = broadcast_shapes(*args) + + @pytest.mark.parametrize( + "args", + [ + ((None,), (None,)), + ((math.nan,), (None,)), + ((1, None, 2, 4), (2, 3, None, 1), (2, None, None, 4)), + ((1, math.nan, 2), (4, 2, 3, math.nan), (4, 2, None, None)), + ((math.nan, 1), (None, 2), (None, 2)), + ], + ) + def test_none(self, args: tuple[tuple[float | None, ...], ...]): + expect = args[-1] + actual = broadcast_shapes(*args[:-1]) + assert actual == expect + + +@pytest.mark.filterwarnings(r"ignore:.*removed in v1.0.0.*:DeprecationWarning") +class TestExpandDims: + def test_single_axis(self, xp: ArrayNamespace): + """Trivial case where xpx.expand_dims doesn't add anything to xp.expand_dims""" + a = xp.asarray(np.reshape(np.arange(2 * 3 * 4 * 5), (2, 3, 4, 5))) + for axis in range(-5, 4): + b = expand_dims(a, axis=axis) + assert_equal(b, xp.expand_dims(a, axis=axis)) + + def test_axis_tuple(self, xp: ArrayNamespace): + a = xp.empty((3, 3, 3)) + assert expand_dims(a, axis=(0, 1, 2)).shape == (1, 1, 1, 3, 3, 3) + assert expand_dims(a, axis=(0, -1, -2)).shape == (1, 3, 3, 3, 1, 1) + assert expand_dims(a, axis=(0, 3, 5)).shape == (1, 3, 3, 1, 3, 1) + assert expand_dims(a, axis=(0, -3, -5)).shape == (1, 1, 3, 1, 3, 3) + + def test_axis_out_of_range(self, xp: ArrayNamespace): + a = xp.empty((2, 3, 4, 5)) + with pytest.raises(IndexError, match="out of bounds"): + _ = expand_dims(a, axis=-6) + with pytest.raises(IndexError, match="out of bounds"): + _ = expand_dims(a, axis=5) + + a = xp.empty((3, 3, 3)) + with pytest.raises(IndexError, match="out of bounds"): + _ = expand_dims(a, axis=(0, -6)) + with pytest.raises(IndexError, match="out of bounds"): + _ = expand_dims(a, axis=(0, 5)) + + def test_repeated_axis(self, xp: ArrayNamespace): + a = xp.empty((3, 3, 3)) + with pytest.raises(ValueError, match="Duplicate dimensions"): + _ = expand_dims(a, axis=(1, 1)) + + def test_positive_negative_repeated(self, xp: ArrayNamespace): + # https://github.com/data-apis/array-api/issues/760#issuecomment-1989449817 + a = xp.empty((2, 3, 4, 5)) + with pytest.raises(ValueError, match="Duplicate dimensions"): + _ = expand_dims(a, axis=(3, -3)) + + def test_device(self, xp: ArrayNamespace, device: Device): + x = xp.asarray([1, 2, 3], device=device) + assert get_device(expand_dims(x, axis=0)) == device + + def test_xp(self, xp: ArrayNamespace): + x = xp.asarray([1, 2, 3]) + y = expand_dims(x, axis=(0, 1, 2), xp=xp) + assert y.shape == (1, 1, 1, 3) + + +class TestPad: + def test_simple(self, xp: ArrayNamespace): + a = xp.asarray([1, 2, 3]) + padded = pad(a, 2) + assert_equal(padded, xp.asarray([0, 0, 1, 2, 3, 0, 0])) + + @pytest.mark.xfail_xp_backend( + Backend.SPARSE, reason="constant_values can only be equal to fill value" + ) + def test_fill_value(self, xp: ArrayNamespace): + a = xp.asarray([1, 2, 3]) + padded = pad(a, 2, constant_values=42) + assert_equal(padded, xp.asarray([42, 42, 1, 2, 3, 42, 42])) + + def test_ndim(self, xp: ArrayNamespace): + a = xp.asarray(np.reshape(np.arange(2 * 3 * 4), (2, 3, 4))) + padded = pad(a, 2) + assert padded.shape == (6, 7, 8) + + def test_mode_not_implemented(self, xp: ArrayNamespace): + a = xp.asarray([1, 2, 3]) + with pytest.raises(NotImplementedError, match="Only `'constant'`"): + _ = pad(a, 2, mode="edge") # type: ignore[arg-type] # pyright: ignore[reportArgumentType] + + def test_device(self, xp: ArrayNamespace, device: Device): + a = xp.asarray(0.0, device=device) + assert get_device(pad(a, 2)) == device + + def test_xp(self, xp: ArrayNamespace): + padded = pad(xp.asarray(0), 1, xp=xp) + assert_equal(padded, xp.asarray(0)) + + def test_tuple_width(self, xp: ArrayNamespace): + a = xp.asarray(np.reshape(np.arange(12), (3, 4))) + padded = pad(a, (1, 0)) + assert padded.shape == (4, 5) + + padded = pad(a, (1, 2)) + assert padded.shape == (6, 7) + + with pytest.raises((ValueError, RuntimeError)): + _ = pad(a, [(1, 2, 3)]) # type: ignore[list-item] # pyright: ignore[reportArgumentType] + + def test_sequence_of_tuples_width(self, xp: ArrayNamespace): + a = xp.asarray(np.reshape(np.arange(12), (3, 4))) + + padded = pad(a, ((1, 0), (0, 2))) + assert padded.shape == (4, 6) + padded = pad(a, ((1, 0), (0, 0))) + assert padded.shape == (4, 4) diff --git a/tests/test_public_api.py b/tests/test_public_api.py new file mode 100644 index 00000000..617cdc27 --- /dev/null +++ b/tests/test_public_api.py @@ -0,0 +1,59 @@ +import inspect + +from array_api_extra import ( + _agnostic, + _at, + _creation, + _elementwise, + _indexing, + _lazy, + _linalg, + _manipulation, + _searching, + _set, + _sorting, + _statistical, + testing, +) + + +def test_all_contains_all_public_functions(): + for module in ( + _at, + _creation, + _elementwise, + _indexing, + _lazy, + _linalg, + _manipulation, + _searching, + _set, + _sorting, + _statistical, + _agnostic._creation, + _agnostic._elementwise, + _agnostic._indexing, + _agnostic._inspection, + _agnostic._linalg, + _agnostic._manipulation, + _agnostic._searching, + _agnostic._set, + _agnostic._sorting, + _agnostic._statistical, + testing._testing, + ): + + def is_function_or_class(member: object): + return inspect.isfunction(member) or inspect.isclass(member) + + public_functions_classes = { + name + for name, obj in inspect.getmembers(module, is_function_or_class) + if not name.startswith("_") and obj.__module__ == module.__name__ + } + missing = sorted(public_functions_classes - set(module.__all__)) + extra = sorted(set(module.__all__) - public_functions_classes) + assert public_functions_classes == set(module.__all__), ( + f"{module.__name__}: Missing from __all__: {missing}\t" + f"Extra in __all__: {extra}" + ) diff --git a/tests/test_searching.py b/tests/test_searching.py new file mode 100644 index 00000000..36298889 --- /dev/null +++ b/tests/test_searching.py @@ -0,0 +1,208 @@ +from collections.abc import Callable +from typing import Any, Literal, cast + +import numpy as np +import pytest + +from array_api_extra import default_dtype +from array_api_extra import searchsorted as xpx_searchsorted +from array_api_extra._agnostic._searching import searchsorted as _funcs_searchsorted +from array_api_extra._lib._backends import Backend +from array_api_extra._lib._compat import ( + array_namespace, + is_jax_namespace, + is_torch_namespace, +) +from array_api_extra._lib._typing import Array, ArrayNamespace +from array_api_extra.testing import assert_equal, lazy_xp_function + +lazy_xp_function(xpx_searchsorted) +lazy_xp_function(_funcs_searchsorted) + + +def _apply_over_batch(*argdefs: tuple[str, int]) -> Any: + """ + Factory for decorator that applies a function over batched arguments. + + Copied (with light simplifications) from `scipy._lib._util`. + + Array arguments may have any number of core dimensions (typically 0, + 1, or 2) and any broadcastable batch shapes. There may be any + number of array outputs of any number of dimensions. Assumptions + right now - which are satisfied by all functions of interest in `linalg` - + are that all array inputs are consecutive keyword or positional arguments, + and that the wrapped function returns either a single array or a tuple of + arrays. It's only as general as it needs to be right now - it can be extended. + + Parameters + ---------- + *argdefs : tuple of (str, int) + Definitions of array arguments: the keyword name of the argument, and + the number of core dimensions. + + Example: + -------- + `linalg.eig` accepts two matrices as the first two arguments `a` and `b`, where + `b` is optional, and returns one array or a tuple of arrays, depending on the + values of other positional or keyword arguments. To generate a wrapper that applies + the function over batches of `a` and optionally `b` : + + >>> _apply_over_batch(('a', 2), ('b', 2)) + """ + names, ndims = list(zip(*argdefs, strict=True)) + n_arrays = len(names) + + def decorator(f: Any) -> Any: + def wrapper( + *args_tuple: Any, + **kwargs: Any, + ) -> Any: + args = list(args_tuple) + + # Ensure all arrays in `arrays`, other arguments in `other_args`/`kwargs` + arrays, other_args = args[:n_arrays], args[n_arrays:] + arrays = cast(list[Array | None], arrays) + for i, name in enumerate(names): + if name in kwargs: + if i + 1 <= len(args): + message = ( + f"{f.__name__}() got multiple values for argument `{name}`." + ) + raise ValueError(message) + arrays.append(kwargs.pop(name)) + + xp = array_namespace(*arrays) + + # Determine core and batch shapes + batch_shapes = [] + core_shapes = [] + for i, (array, ndim) in enumerate(zip(arrays, ndims, strict=True)): + array = None if array is None else xp.asarray(array) # noqa: PLW2901 + shape = () if array is None else array.shape + arrays[i] = array + batch_shapes.append(shape[:-ndim] if ndim > 0 else shape) + core_shapes.append(shape[-ndim:] if ndim > 0 else ()) + + # Early exit if call is not batched + if not any(batch_shapes): + return f(*arrays, *other_args, **kwargs) + + # Determine broadcasted batch shape + batch_shape = np.broadcast_shapes(*batch_shapes) # Gives OK error message + + # Broadcast arrays to appropriate shape + for i, (array, core_shape) in enumerate( + zip(arrays, core_shapes, strict=True) + ): + if array is None: + continue + arrays[i] = xp.broadcast_to(array, batch_shape + core_shape) + + # Main loop + results = [] + for index in np.ndindex(batch_shape): + result = f( + *( + (array[index] if array is not None else None) + for array in arrays + ), + *other_args, + **kwargs, + ) + # Assume `result` is either a tuple or single array. This is easily + # generalized by allowing the contributor to pass an `unpack_result` + # callable to the decorator factory. + result = (result,) if not isinstance(result, tuple) else result + results.append(result) + results = list(zip(*results, strict=True)) + + # Reshape results + for i, result in enumerate(results): + result = xp.stack(result) # noqa: PLW2901 + core_shape = result.shape[1:] + results[i] = xp.reshape(result, batch_shape + core_shape) + + # Assume `result` should be a single array if there is only one element or + # a `tuple` otherwise. This is easily generalized by allowing the + # contributor to pass an `pack_result` callable to the decorator factory. + return results[0] if len(results) == 1 else results + + return wrapper + + return decorator + + +@_apply_over_batch(("a", 1), ("v", 1)) # type: ignore[untyped-decorator] +def xp_searchsorted( + a: Array, + v: Array, + side: Literal["left", "right"], + xp: ArrayNamespace, +) -> Array: + return xp.searchsorted(a, v, side=side) + + +@pytest.mark.skip_xp_backend(Backend.DASK, reason="no take_along_axis") +@pytest.mark.skip_xp_backend(Backend.SPARSE, reason="no searchsorted") +class TestSearchsorted: + def test_input_validation(self, xp: ArrayNamespace): + message = "`side` must be either 'left' or 'right'." + with pytest.raises(ValueError, match=message): + _ = xpx_searchsorted(xp.asarray([1, 2]), xp.asarray([1, 2]), side="center") # type: ignore[arg-type] # pyright: ignore[reportArgumentType] + + @pytest.mark.parametrize("side", ["left", "right"]) + @pytest.mark.parametrize("ties", [False, True]) + @pytest.mark.parametrize( + "shape", [0, 1, 2, 10, 11, 1000, 10001, (2, 0), (0, 2), (2, 10), (2, 3, 11)] + ) + @pytest.mark.parametrize("nans_x", [False, True]) + @pytest.mark.parametrize("infs_x", [False, True]) + @pytest.mark.parametrize("searchsorted", [xpx_searchsorted, _funcs_searchsorted]) + def test_nd( + self, + side: Literal["left", "right"], + ties: bool, + shape: int | tuple[int], + nans_x: bool, + infs_x: bool, + xp: ArrayNamespace, + searchsorted: Callable[..., Array], + ): + if nans_x and is_jax_namespace(xp): + pytest.xfail("https://github.com/jax-ml/jax/issues/39887") + if nans_x and is_torch_namespace(xp) and searchsorted == xpx_searchsorted: + pytest.skip("torch sorts NaNs differently") + if isinstance(shape, tuple) and searchsorted == _funcs_searchsorted: + message = ( + "Redundant; `xpx_searchsorted` delegates to " + "`_funcs_searchsorted` for multidimensional input." + ) + pytest.skip(message) + rng = np.random.default_rng(945298725498274853) + x = rng.integers(5, size=shape) if ties else rng.random(shape) + # float32 is to accommodate JAX - nextafter with `float64` is too small? + x = np.asarray(x, dtype=np.float32) # type:ignore[assignment] + xr = np.nextafter(x, np.inf) + xl = np.nextafter(x, -np.inf) + x_ = np.asarray([-np.inf, np.inf, np.nan]) + x_ = np.broadcast_to(x_, (*x.shape[:-1], 3)) + y = rng.permuted(np.concatenate((xl, x, xr, x_), axis=-1), axis=-1) + if nans_x: + mask = rng.random(shape) < 0.1 + x[mask] = np.nan + if infs_x: + mask = rng.random(shape) < 0.1 + x[mask] = -np.inf + mask = rng.random(shape) > 0.9 + x[mask] = np.inf + x = np.sort(x, axis=-1) # type:ignore[assignment] + x, y = np.asarray(x, dtype=np.float64), np.asarray(y, dtype=np.float64) + xp_default_int = default_dtype(xp, kind="integral") + if x.size == 0 and x.ndim > 0 and x.shape[-1] != 0: + ref = xp.empty((*x.shape[:-1], y.shape[-1]), dtype=xp_default_int) + else: + ref = xp_searchsorted(x, y, side=side, xp=np) + ref = xp.asarray(ref, dtype=xp_default_int) + x, y = xp.asarray(x.copy()), xp.asarray(y.copy()) + res = searchsorted(x, y, side=side, xp=xp) + assert_equal(res, ref) diff --git a/tests/test_set.py b/tests/test_set.py new file mode 100644 index 00000000..3048b23a --- /dev/null +++ b/tests/test_set.py @@ -0,0 +1,261 @@ +import pytest + +from array_api_extra import _agnostic, isin, nunique, setdiff1d, union1d +from array_api_extra._lib._backends import NUMPY_VERSION, Backend +from array_api_extra._lib._compat import device as get_device +from array_api_extra._lib._typing import Array, ArrayNamespace, Device +from array_api_extra.testing import assert_equal, lazy_xp_function + +lazy_xp_function(isin) +lazy_xp_function(nunique) +# FIXME calls in1d which calls xp.unique_values without size +lazy_xp_function(setdiff1d, jax_jit=False) +lazy_xp_function(union1d, jax_jit=False) + + +class TestNUnique: + @pytest.mark.skip_xp_backend( + Backend.ARRAY_API_STRICT, reason="array-agnostic fallback" + ) + @pytest.mark.skip_xp_backend( + Backend.ARRAY_API_STRICTEST, reason="array-agnostic fallback" + ) + @pytest.mark.skip_xp_backend(Backend.DASK, reason="array-agnostic fallback") + @pytest.mark.skip_xp_backend(Backend.SPARSE, reason="array-agnostic fallback") + def test_delegates( + self, + xp: ArrayNamespace, + monkeypatch: pytest.MonkeyPatch, + ): + def fallback(*_args: object, **_kwargs: object) -> Array: + msg = "array-agnostic fallback should not be used" + raise AssertionError(msg) + + monkeypatch.setattr(_agnostic._set, "nunique", fallback) + a = xp.asarray([1, 1, 2]) + assert_equal(nunique(a), xp.asarray(2)) + + def test_simple(self, xp: ArrayNamespace): + a = xp.asarray([[1, 1], [0, 2], [2, 2]]) + assert_equal(nunique(a), xp.asarray(3)) + + def test_empty(self, xp: ArrayNamespace): + a = xp.asarray([]) + assert_equal(nunique(a), xp.asarray(0)) + + def test_size1(self, xp: ArrayNamespace): + a = xp.asarray([123]) + assert_equal(nunique(a), xp.asarray(1)) + + def test_all_equal(self, xp: ArrayNamespace): + a = xp.asarray([123, 123, 123]) + assert_equal(nunique(a), xp.asarray(1)) + + @pytest.mark.xfail_xp_backend(Backend.DASK, reason="No equal_nan kwarg in unique") + def test_nan(self, xp: ArrayNamespace, library: Backend): + if library.like(Backend.NUMPY) and NUMPY_VERSION < (1, 24): + pytest.xfail("NumPy <1.24 has no equal_nan kwarg in unique") + + # Each NaN is counted separately + a = xp.asarray([xp.nan, 123.0, xp.nan]) + assert_equal(nunique(a), xp.asarray(3)) + + @pytest.mark.parametrize("size", [0, 1, 2]) + def test_device(self, xp: ArrayNamespace, device: Device, size: int): + a = xp.asarray([0.0] * size, device=device) + assert get_device(nunique(a)) == device + + def test_xp(self, xp: ArrayNamespace): + a = xp.asarray([[1, 1], [0, 2], [2, 2]]) + assert_equal(nunique(a, xp=xp), xp.asarray(3)) + + +assume_unique = pytest.mark.parametrize( + "assume_unique", + [ + True, + pytest.param( + False, + marks=pytest.mark.xfail_xp_backend( + Backend.DASK, reason="NaN-shaped arrays" + ), + ), + ], +) + + +@pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no argsort") +@pytest.mark.skip_xp_backend(Backend.ARRAY_API_STRICTEST, reason="no unique_values") +class TestSetDiff1D: + @pytest.mark.xfail_xp_backend(Backend.DASK, reason="NaN-shaped arrays") + @pytest.mark.xfail_xp_backend( + Backend.TORCH, reason="index_select not implemented for uint32" + ) + @pytest.mark.xfail_xp_backend( + Backend.TORCH_GPU, reason="index_select not implemented for uint32" + ) + def test_setdiff1d(self, xp: ArrayNamespace): + x1 = xp.asarray([6, 5, 4, 7, 1, 2, 7, 4]) + x2 = xp.asarray([2, 4, 3, 3, 2, 1, 5]) + + expected = xp.asarray([6, 7]) + actual = setdiff1d(x1, x2) + assert_equal(actual, expected) + + x1 = xp.arange(21) + x2 = xp.arange(19) + expected = xp.asarray([19, 20]) + actual = setdiff1d(x1, x2) + assert_equal(actual, expected) + + assert_equal(setdiff1d(xp.empty(0), xp.empty(0)), xp.empty(0)) + x1 = xp.empty(0, dtype=xp.uint32) + x2 = x1 + assert xp.isdtype(setdiff1d(x1, x2).dtype, xp.uint32) + + def test_assume_unique(self, xp: ArrayNamespace): + x1 = xp.asarray([3, 2, 1]) + x2 = xp.asarray([7, 5, 2]) + expected = xp.asarray([3, 1]) + actual = setdiff1d(x1, x2, assume_unique=True) + assert_equal(actual, expected) + + @assume_unique + @pytest.mark.parametrize("shape1", [(), (1,), (1, 1)]) + @pytest.mark.parametrize("shape2", [(), (1,), (1, 1)]) + def test_shapes( + self, + assume_unique: bool, + shape1: tuple[int, ...], + shape2: tuple[int, ...], + xp: ArrayNamespace, + ): + x1 = xp.zeros(shape1) + x2 = xp.zeros(shape2) + + actual = setdiff1d(x1, x2, assume_unique=assume_unique) + assert_equal(actual, xp.empty((0,))) + + @assume_unique + @pytest.mark.skip_xp_backend(Backend.NUMPY_READONLY, reason="xp=xp") + def test_python_scalar(self, xp: ArrayNamespace, assume_unique: bool): + # Test no dtype promotion to xp.asarray(x2); use x1.dtype + x1 = xp.asarray([3, 1, 2], dtype=xp.int16) + x2 = 3 + actual = setdiff1d(x1, x2, assume_unique=assume_unique) + assert_equal(actual, xp.asarray([1, 2], dtype=xp.int16)) + + actual = setdiff1d(x2, x1, assume_unique=assume_unique) + assert_equal(actual, xp.asarray([], dtype=xp.int16)) + + assert_equal( + setdiff1d(0, 0, assume_unique=assume_unique, xp=xp), + xp.asarray([0])[:0], # Default int dtype for backend + ) + + @pytest.mark.parametrize("assume_unique", [True, False]) + def test_all_python_scalars(self, assume_unique: bool): + with pytest.raises(TypeError, match=r"array_namespace requires .* array input"): + _ = setdiff1d(0, 0, assume_unique=assume_unique) + + @assume_unique + @pytest.mark.skip_xp_backend( + Backend.TORCH, reason="device='meta' does not support unknown shapes" + ) + def test_device(self, xp: ArrayNamespace, device: Device, assume_unique: bool): + x1 = xp.asarray([3, 8, 20], device=device) + x2 = xp.asarray([2, 3, 4], device=device) + assert get_device(setdiff1d(x1, x2, assume_unique=assume_unique)) == device + + @pytest.mark.skip_xp_backend(Backend.NUMPY_READONLY, reason="xp=xp") + def test_xp(self, xp: ArrayNamespace): + x1 = xp.asarray([3, 8, 20]) + x2 = xp.asarray([2, 3, 4]) + expected = xp.asarray([8, 20]) + actual = setdiff1d(x1, x2, assume_unique=True, xp=xp) + assert_equal(actual, expected) + + +@pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no unique_inverse") +class TestIsIn: + def test_simple(self, xp: ArrayNamespace, library: Backend): + if library.like(Backend.NUMPY) and NUMPY_VERSION < (1, 24): + pytest.xfail("NumPy <1.24 has no kind kwarg in isin") + + b = xp.asarray([1, 2, 3, 4]) + + # `a` with 1 dimension + a = xp.asarray([1, 3, 6, 10]) + expected = xp.asarray([True, True, False, False]) + res = isin(a, b) + assert_equal(res, expected) + + # `a` with 2 dimensions + a = xp.asarray([[0, 2], [4, 6]]) + expected = xp.asarray([[False, True], [True, False]]) + res = isin(a, b) + assert_equal(res, expected) + + def test_device(self, xp: ArrayNamespace, device: Device, library: Backend): + if library.like(Backend.NUMPY) and NUMPY_VERSION < (1, 24): + pytest.xfail("NumPy <1.24 has no kind kwarg in isin") + + a = xp.asarray([1, 3, 6], device=device) + b = xp.asarray([1, 2, 3], device=device) + assert get_device(isin(a, b)) == device + + def test_assume_unique_and_invert( + self, xp: ArrayNamespace, device: Device, library: Backend + ): + if library.like(Backend.NUMPY) and NUMPY_VERSION < (1, 24): + pytest.xfail("NumPy <1.24 has no kind kwarg in isin") + + a = xp.asarray([0, 3, 6, 10], device=device) + b = xp.asarray([1, 2, 3, 10], device=device) + expected = xp.asarray([True, False, True, False], device=device) + res = isin(a, b, assume_unique=True, invert=True) + assert get_device(res) == device + assert_equal(res, expected) + + def test_kind(self, xp: ArrayNamespace, library: Backend): + if library.like(Backend.NUMPY) and NUMPY_VERSION < (1, 24): + pytest.xfail("NumPy <1.24 has no kind kwarg in isin") + + a = xp.asarray([0, 3, 6, 10]) + b = xp.asarray([1, 2, 3, 10]) + expected = xp.asarray([False, True, False, True]) + res = isin(a, b, kind="sort") + assert_equal(res, expected) + + +@pytest.mark.skip_xp_backend( + Backend.ARRAY_API_STRICTEST, + reason="data_dependent_shapes flag for unique_values is disabled", +) +class TestUnion1d: + def test_simple(self, xp: ArrayNamespace): + a = xp.asarray([-1, 1, 0]) + b = xp.asarray([2, -2, 0]) + expected = xp.asarray([-2, -1, 0, 1, 2]) + res = union1d(a, b) + assert_equal(res, expected) + + def test_2d(self, xp: ArrayNamespace): + a = xp.asarray([[-1, 1, 0], [1, 2, 0]]) + b = xp.asarray([[1, 0, 1], [-2, -1, 0]]) + expected = xp.asarray([-2, -1, 0, 1, 2]) + res = union1d(a, b) + assert_equal(res, expected) + + def test_3d(self, xp: ArrayNamespace): + a = xp.asarray([[[-1, 0], [1, 2]], [[-1, 0], [1, 2]]]) + b = xp.asarray([[[0, 1], [-1, 2]], [[1, -2], [0, 2]]]) + expected = xp.asarray([-2, -1, 0, 1, 2]) + res = union1d(a, b) + assert_equal(res, expected) + + @pytest.mark.skip_xp_backend(Backend.TORCH, reason="materialize 'meta' device") + def test_device(self, xp: ArrayNamespace, device: Device): + a = xp.asarray([-1, 1, 0], device=device) + b = xp.asarray([2, -2, 0], device=device) + assert get_device(union1d(a, b)) == device diff --git a/tests/test_sorting.py b/tests/test_sorting.py new file mode 100644 index 00000000..37d91b9e --- /dev/null +++ b/tests/test_sorting.py @@ -0,0 +1,153 @@ +import numpy as np +import pytest +from typing_extensions import override + +from array_api_extra import argpartition, partition +from array_api_extra._lib._backends import Backend +from array_api_extra._lib._typing import Array, ArrayNamespace +from array_api_extra.testing import lazy_xp_function + +lazy_xp_function(argpartition) +lazy_xp_function(partition) + + +class TestPartition: + @classmethod + def _assert_valid_partition( + cls, + x_np: np.ndarray | None, + k: int, + y: Array, + xp: ArrayNamespace, + axis: int | None = -1, + ): + """ + x_np : input array + k : int + y : output array returned by the partition function to test + """ + if x_np is not None: + assert y.shape == np.partition(x_np, k, axis=axis).shape + if y.ndim != 1 and axis == 0: + assert isinstance(y.shape[1], int) + for i in range(y.shape[1]): + cls._assert_valid_partition(None, k, y[:, i, ...], xp, axis=0) + elif y.ndim != 1: + assert axis is not None + axis = axis - 1 if axis != -1 else -1 + assert isinstance(y.shape[0], int) + for i in range(y.shape[0]): + cls._assert_valid_partition(None, k, y[i, ...], xp, axis=axis) + else: + if k > 0: + assert xp.max(y[:k]) <= y[k] + assert y[k] <= xp.min(y[k:]) + + @classmethod + def _partition( + cls, x: np.ndarray, k: int, xp: ArrayNamespace, axis: int | None = -1 + ): + return partition(xp.asarray(x), k, axis=axis) + + def _test_1d(self, xp: ArrayNamespace): + rng = np.random.default_rng() + for n in [2, 3, 4, 5, 7, 10, 20, 50, 100, 1_000]: + k = int(rng.integers(n)) + x1 = rng.integers(n, size=n) + y = self._partition(x1, k, xp) + self._assert_valid_partition(x1, k, y, xp) + x2 = rng.random(n) + y = self._partition(x2, k, xp) + self._assert_valid_partition(x2, k, y, xp) + + def _test_nd(self, xp: ArrayNamespace, ndim: int): + rng = np.random.default_rng() + + for n in [2, 3, 5, 10, 20, 100]: + base_shape = [int(v) for v in rng.integers(1, 4, size=ndim)] + k = int(rng.integers(n)) + + for i in range(ndim): + shape = base_shape[:] + shape[i] = n + x = rng.integers(n, size=tuple(shape)) + y = self._partition(x, k, xp, axis=i) + self._assert_valid_partition(x, k, y, xp, axis=i) + + z = rng.random(tuple(base_shape)) + k = int(rng.integers(z.size)) + y = self._partition(z, k, xp, axis=None) + self._assert_valid_partition(z, k, y, xp, axis=None) + + def _test_input_validation(self, xp: ArrayNamespace): + with pytest.raises(TypeError): + _ = self._partition(np.asarray(1), 1, xp) + with pytest.raises(ValueError, match="out of bounds"): + _ = self._partition(np.asarray([1, 2]), 3, xp) + + def test_1d(self, xp: ArrayNamespace): + self._test_1d(xp) + + @pytest.mark.parametrize("ndim", [2, 3, 4]) + def test_nd(self, xp: ArrayNamespace, ndim: int): + self._test_nd(xp, ndim) + + def test_input_validation(self, xp: ArrayNamespace): + self._test_input_validation(xp) + + +@pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no argsort") +class TestArgpartition(TestPartition): + @classmethod + @override + def _partition( + cls, x: np.ndarray, k: int, xp: ArrayNamespace, axis: int | None = -1 + ): + arr = xp.asarray(x) + indices = argpartition(arr, k, axis=axis) + if axis is None: + arr = xp.reshape(arr, shape=(-1,)) + return arr[indices] + if arr.ndim == 1: + return arr[indices] + return cls._take_along_axis(arr, indices, axis=axis, xp=xp) + + @classmethod + def _take_along_axis( + cls, arr: Array, indices: Array, axis: int, xp: ArrayNamespace + ): + if hasattr(xp, "take_along_axis"): + return xp.take_along_axis(arr, indices, axis=axis) + if arr.ndim == 1: + return arr[indices] + if axis == 0: + assert isinstance(arr.shape[1], int) + arrs = [] + for i in range(arr.shape[1]): + arrs.append( + cls._take_along_axis( + arr[:, i, ...], indices[:, i, ...], axis=0, xp=xp + ) + ) + return xp.stack(arrs, axis=1) + axis = axis - 1 if axis != -1 else -1 + assert isinstance(arr.shape[0], int) + arrs = [] + for i in range(arr.shape[0]): + arrs.append( + cls._take_along_axis(arr[i, ...], indices[i, ...], axis=axis, xp=xp) + ) + return xp.stack(arrs, axis=0) + + @override + def test_1d(self, xp: ArrayNamespace): + self._test_1d(xp) + + @pytest.mark.parametrize("ndim", [2, 3, 4]) + @override + def test_nd(self, xp: ArrayNamespace, ndim: int): + self._test_nd(xp, ndim) + + @override + def test_input_validation(self, xp: ArrayNamespace): + self._test_input_validation(xp) diff --git a/tests/test_statistical.py b/tests/test_statistical.py new file mode 100644 index 00000000..91319645 --- /dev/null +++ b/tests/test_statistical.py @@ -0,0 +1,508 @@ +import math +import warnings +from typing import Any, cast + +import numpy as np +import pytest + +from array_api_extra import cov, nanmax, nanmin, nansum +from array_api_extra._lib._backends import Backend +from array_api_extra._lib._compat import array_namespace +from array_api_extra._lib._compat import device as get_device +from array_api_extra._lib._typing import Array, ArrayNamespace, Device +from array_api_extra.testing import assert_close, assert_equal, lazy_xp_function + +lazy_xp_function(cov) +lazy_xp_function(nansum) + + +class TestCov: + def test_basic(self, xp: ArrayNamespace): + assert_close( + cov(xp.asarray([[0, 2], [1, 1], [2, 0]], dtype=xp.float64).T), + xp.asarray([[1.0, -1.0], [-1.0, 1.0]], dtype=xp.float64), + ) + + def test_complex(self, xp: ArrayNamespace): + actual = cov(xp.asarray([[1, 2, 3], [1j, 2j, 3j]], dtype=xp.complex128)) + expect = xp.asarray([[1.0, -1.0j], [1.0j, 1.0]], dtype=xp.complex128) + assert_close(actual, expect) + + def test_complex_with_weights(self, xp: ArrayNamespace): + m = np.asarray( + [[1 + 1j, 2 + 2j, 4 + 1j], [3 - 1j, 5 + 2j, 7 + 0j]], + dtype=np.complex128, + ) + weights = np.asarray([1.0, 2.0, 1.0]) + correction = 0.5 # Force the generic implementation. + + weight_sum = weights.sum() + avg = (m * weights).sum(axis=-1, keepdims=True) / weight_sum + centered = m - avg + normalizer = weight_sum - correction * (weights**2).sum() / weight_sum + expected = (centered * weights) @ centered.conj().T / normalizer + + actual = cov( + xp.asarray(m), + correction=correction, + aweights=xp.asarray(weights), + ) + assert_close(actual, xp.asarray(expected)) + + def test_empty(self, xp: ArrayNamespace): + with warnings.catch_warnings(record=True): + warnings.simplefilter("always", RuntimeWarning) + warnings.simplefilter("always", UserWarning) + assert_equal( + cov(xp.asarray([], dtype=xp.float64)), + xp.asarray(xp.nan, dtype=xp.float64), + ) + assert_equal( + cov(xp.reshape(xp.asarray([], dtype=xp.float64), (0, 2))), + xp.reshape(xp.asarray([], dtype=xp.float64), (0, 0)), + ) + assert_equal( + cov(xp.reshape(xp.asarray([], dtype=xp.float64), (2, 0))), + xp.asarray([[xp.nan, xp.nan], [xp.nan, xp.nan]], dtype=xp.float64), + ) + + def test_combination(self, xp: ArrayNamespace): + x = xp.asarray([-2.1, -1, 4.3], dtype=xp.float64) + y = xp.asarray([3, 1.1, 0.12], dtype=xp.float64) + X = xp.stack((x, y), axis=0) + desired = xp.asarray([[11.71, -4.286], [-4.286, 2.144133]], dtype=xp.float64) + assert_close(cov(X), desired, rtol=1e-6) + assert_close(cov(x), xp.asarray(11.71, dtype=xp.float64)) + assert_close(cov(y), xp.asarray(2.144133, dtype=xp.float64), rtol=1e-6) + + @pytest.mark.xfail_xp_backend( + Backend.TORCH, reason="torch.cov does not support tensors on meta device" + ) + def test_device(self, xp: ArrayNamespace, device: Device): + x = xp.asarray([1, 2, 3], device=device) + assert get_device(cov(x)) == device + + @pytest.mark.skip_xp_backend(Backend.NUMPY_READONLY, reason="xp=xp") + def test_xp(self, xp: ArrayNamespace): + assert_close( + cov( + xp.asarray([[0.0, 2.0], [1.0, 1.0], [2.0, 0.0]], dtype=xp.float64).T, + xp=xp, + ), + xp.asarray([[1.0, -1.0], [-1.0, 1.0]], dtype=xp.float64), + ) + + def test_batch(self, xp: ArrayNamespace): + rng = np.random.default_rng(8847643423) + batch_shape = (3, 4) + n_var, n_obs = 3, 20 + m = rng.random((*batch_shape, n_var, n_obs)) + res = cov(xp.asarray(m)) + ref_list = [np.cov(m_) for m_ in np.reshape(m, (-1, n_var, n_obs))] + ref = np.reshape(np.stack(ref_list), (*batch_shape, n_var, n_var)) + assert_close(res, xp.asarray(ref)) + + @pytest.mark.parametrize("bias", [True, False, 0, 1]) + def test_bias(self, xp: ArrayNamespace, bias: bool): + # `bias` maps to `correction`: bias=True -> correction=0, bias=False -> 1. + x = np.array([-2.1, -1, 4.3]) + y = np.array([3, 1.1, 0.12]) + X = np.stack((x, y), axis=0) + ref = np.cov(X, bias=bias) + assert_close( + cov(xp.asarray(X, dtype=xp.float64), correction=0 if bias else 1), + xp.asarray(ref, dtype=xp.float64), + rtol=1e-6, + ) + + @pytest.mark.parametrize("bias", [True, False, 0, 1]) + def test_bias_batch(self, xp: ArrayNamespace, bias: bool): + rng = np.random.default_rng(8847643423) + batch_shape = (3, 4) + n_var, n_obs = 3, 20 + m = rng.random((*batch_shape, n_var, n_obs)) + res = cov(xp.asarray(m), correction=0 if bias else 1) + ref_list = [np.cov(m_, bias=bias) for m_ in np.reshape(m, (-1, n_var, n_obs))] + ref = np.reshape(np.stack(ref_list), (*batch_shape, n_var, n_var)) + assert_close(res, xp.asarray(ref)) + + def test_correction(self, xp: ArrayNamespace): + rng = np.random.default_rng(20260417) + m = rng.random((3, 20)) + for correction in (0, 1, 2): + ref = np.cov(m, ddof=correction) + res = cov(xp.asarray(m), correction=correction) + assert_close(res, xp.asarray(ref)) + + def test_correction_float(self, xp: ArrayNamespace): + # Float correction: reference computed by hand (numpy.cov rejects + # non-integer ddof; our generic path supports it). + rng = np.random.default_rng(20260417) + m = rng.random((3, 20)) + n = m.shape[-1] + centered = m - m.mean(axis=-1, keepdims=True) + ref = centered @ centered.T / (n - 1.5) + res = cov(xp.asarray(m), correction=1.5) + assert_close(res, xp.asarray(ref)) + + def test_axis(self, xp: ArrayNamespace): + rng = np.random.default_rng(20260417) + m = rng.random((20, 3)) # observations on axis 0 + ref = np.cov(m, rowvar=False) + res = cov(xp.asarray(m), axis=0) + assert_close(res, xp.asarray(ref)) + res_neg = cov(xp.asarray(m), axis=-2) + assert_close(res_neg, xp.asarray(ref)) + + def test_frequency_weights(self, xp: ArrayNamespace): + rng = np.random.default_rng(20260417) + m = rng.random((3, 10)) + fw = np.asarray([1, 2, 1, 3, 1, 2, 1, 1, 2, 1], dtype=np.int64) + ref = np.cov(m, fweights=fw) + res = cov(xp.asarray(m), fweights=xp.asarray(fw)) + assert_close(res, xp.asarray(ref)) + + def test_weights(self, xp: ArrayNamespace): + rng = np.random.default_rng(20260417) + m = rng.random((3, 10)) + aw = rng.random(10) + ref = np.cov(m, aweights=aw) + res = cov(xp.asarray(m), aweights=xp.asarray(aw)) + assert_close(res, xp.asarray(ref)) + + def test_both_weights(self, xp: ArrayNamespace): + rng = np.random.default_rng(20260417) + m = rng.random((3, 10)) + fw = np.asarray([1, 2, 1, 3, 1, 2, 1, 1, 2, 1], dtype=np.int64) + aw = rng.random(10) + for correction in (0, 1, 2): + ref = np.cov(m, ddof=correction, fweights=fw, aweights=aw) + res = cov( + xp.asarray(m), + correction=correction, + fweights=xp.asarray(fw), + aweights=xp.asarray(aw), + ) + assert_close(res, xp.asarray(ref)) + + def test_batch_with_weights(self, xp: ArrayNamespace): + rng = np.random.default_rng(20260417) + batch_shape = (2, 3) + n_var, n_obs = 3, 15 + m = rng.random((*batch_shape, n_var, n_obs)) + aw = rng.random(n_obs) + res = cov(xp.asarray(m), aweights=xp.asarray(aw)) + ref_list = [np.cov(m_, aweights=aw) for m_ in np.reshape(m, (-1, n_var, n_obs))] + ref = np.reshape(np.stack(ref_list), (*batch_shape, n_var, n_var)) + assert_close(res, xp.asarray(ref)) + + def test_axis_with_weights(self, xp: ArrayNamespace): + # axis=-2 (observations on first of 2D) combined with weights: + # verifies that moveaxis and weight alignment cooperate. + rng = np.random.default_rng(20260417) + m = rng.random((15, 3)) # observations on axis 0 + aw = rng.random(15) + fw = np.asarray([1, 2, 1, 3, 1, 2, 1, 1, 2, 1, 1, 1, 2, 1, 1], dtype=np.int64) + ref = np.cov(m, rowvar=False, fweights=fw, aweights=aw) + res = cov( + xp.asarray(m), + axis=-2, + fweights=xp.asarray(fw), + aweights=xp.asarray(aw), + ) + assert_close(res, xp.asarray(ref)) + + def test_axis_out_of_bounds(self, xp: ArrayNamespace): + m = xp.asarray([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + with pytest.raises(IndexError): + _ = cov(m, axis=5) + + def test_weights_wrong_ndim(self, xp: ArrayNamespace): + m = xp.asarray([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + w2d = xp.asarray([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]) + # Non-integer correction forces the generic path where the + # validation lives; native backends raise for the same reason. + with pytest.raises((ValueError, TypeError)): + _ = cov(m, correction=0.5, fweights=w2d) + with pytest.raises((ValueError, TypeError)): + _ = cov(m, correction=0.5, aweights=w2d) + + def test_weights_wrong_length(self, xp: ArrayNamespace): + m = xp.asarray([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + w_bad = xp.asarray([1.0, 1.0]) # expected length 3 + with pytest.raises((ValueError, RuntimeError)): + _ = cov(m, correction=0.5, fweights=w_bad) + with pytest.raises((ValueError, RuntimeError)): + _ = cov(m, correction=0.5, aweights=w_bad) + + def test_weights_unknown_length(self, da: ArrayNamespace): + m_np = np.asarray([[1.0, 2.0, 3.0], [4.0, 6.0, 8.0]]) + weights_np = np.asarray([1.0, 2.0, 3.0]) + keep_np = np.asarray([True, False, True]) + + keep = da.asarray(keep_np) + m = da.asarray(m_np)[:, keep] + weights = da.asarray(weights_np)[keep] + assert math.isnan(m.shape[-1]) + assert math.isnan(weights.shape[0]) + + actual = cov(m, aweights=weights) + desired = np.cov(m_np[:, keep_np], aweights=weights_np[keep_np]) + assert_close(actual, da.asarray(desired)) + + def test_weights_dof_warning_eager(self): + xp = array_namespace(cast(Array, cast(object, np.empty(0)))) + m = xp.asarray([[1.0, 2.0], [3.0, 4.0]]) + weights = xp.asarray([1.0, 1.0]) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + _ = cov(m, correction=2.5, aweights=weights) + assert any( + isinstance(warning.message, RuntimeWarning) + and "Degrees of freedom <= 0" in str(warning.message) + for warning in caught + ) + + def test_torch_autograd(self, torch: ArrayNamespace): + # The batched (generic) path must not detach gradients or mutate the + # input tensor in place, as `xp.asarray` does on torch. + xp = torch + rng = np.random.default_rng(20260417) + m = xp.asarray(rng.random((4, 3, 20)), dtype=xp.float64) + m.requires_grad_(True) + m_before = m.detach().clone() + # cov returns the array-api `Array` type; at runtime it is a torch + # tensor, so cast to access autograd attributes without type errors. + c = cast(Any, cov(m)) # batched -> generic path + assert c.requires_grad + assert m.requires_grad # input tensor not mutated + assert_equal(m.detach(), m_before) + c.sum().backward() + assert m.grad is not None + assert bool(xp.all(xp.isfinite(m.grad))) + + +class TestNanMin: + def test_simple(self, xp: ArrayNamespace): + a = xp.asarray([[1, 2], [3, xp.nan]]) + + # with the default `axis=None` a single scalar is returned + res = nanmin(a) + expected = 1.0 + assert res == expected + + res = nanmin(a, axis=0) + expected = xp.asarray([1.0, 2.0]) + assert_equal(res, expected) + + res = nanmin(a, axis=1) + expected = xp.asarray([1.0, 3.0]) + assert_equal(res, expected) + + def test_bigger(self, xp: ArrayNamespace): + a = xp.asarray( + [ + [1, xp.nan, 4, 5], + [xp.nan, -2, xp.nan, -4], + [2, 1, 3, xp.nan], + ] + ) + + res = nanmin(a, axis=0) + expected = xp.asarray([1.0, -2.0, 3.0, -4.0]) + assert_equal(res, expected) + + res = nanmin(a, axis=1) + expected = xp.asarray([1.0, -4.0, 1.0]) + assert_equal(res, expected) + + def test_with_infinity(self, xp: ArrayNamespace): + a = xp.asarray([0.1, 1.0, xp.nan, xp.inf]) + res = nanmin(a) + expected = 0.1 + assert res == expected + + a = xp.asarray([0.1, 1.0, xp.nan, -xp.inf]) + res = nanmin(a) + expected = -xp.inf + assert res == expected + + def test_scalar(self, xp: ArrayNamespace): + a = xp.asarray(1.0) + assert nanmin(a) == 1.0 + + @pytest.mark.filterwarnings("ignore:.*All-NaN slice*.:RuntimeWarning") + def test_all_nan_slice_2d(self, xp: ArrayNamespace): + a = xp.asarray( + [ + [xp.nan, 5.0], + [xp.nan, 2.0], + ] + ) + + res = nanmin(a, axis=0, xp=xp) + expected = xp.asarray([xp.nan, 2.0]) + assert_equal(res, expected) + + @pytest.mark.skip_xp_backend( + Backend.TORCH, reason="torch.nanmin does not support tensors on meta device" + ) + @pytest.mark.parametrize("axis", [None, 0, 1]) + def test_device(self, axis: int | None, xp: ArrayNamespace, device: Device): + a = xp.asarray([[4, xp.nan, 1], [2, 5, xp.nan]], device=device) + res = nanmin(a, axis=axis) + assert get_device(res) == device + + @pytest.mark.parametrize( + ("axis", "expected_list"), [(0, [2.0, 3.0, 1.0]), (1, [1.0, 2.0])] + ) + def test_xp(self, axis: int | None, expected_list: list[float], xp: ArrayNamespace): + a = xp.asarray([[4, xp.nan, 1], [2, 3, xp.nan]]) + res = nanmin(a, axis=axis, xp=xp) + expected = xp.asarray(expected_list) + assert_equal(res, expected) + + +class TestNanMax: + def test_simple(self, xp: ArrayNamespace): + a = xp.asarray([[5, 3], [6, xp.nan]]) + + # with the default `axis=None` a single scalar is returned + res = nanmax(a) + expected = 6.0 + assert res == expected + + res = nanmax(a, axis=0) + expected = xp.asarray([6.0, 3.0]) + assert_equal(res, expected) + + res = nanmax(a, axis=1) + expected = xp.asarray([5.0, 6.0]) + assert_equal(res, expected) + + def test_bigger(self, xp: ArrayNamespace): + a = xp.asarray( + [ + [1, xp.nan, 4, 5], + [xp.nan, 2, xp.nan, 4], + [6, 1, 3, xp.nan], + ] + ) + + res = nanmax(a, axis=0) + expected = xp.asarray([6.0, 2.0, 4.0, 5.0]) + assert_equal(res, expected) + + res = nanmax(a, axis=1) + expected = xp.asarray([5.0, 4.0, 6.0]) + assert_equal(res, expected) + + def test_with_infinity(self, xp: ArrayNamespace): + a = xp.asarray([0.1, 5.0, xp.nan, -xp.inf]) + res = nanmax(a) + expected = 5.0 + assert res == expected + + a = xp.asarray([3.0, 10.0, xp.nan, xp.inf]) + res = nanmax(a) + expected = xp.inf + assert res == expected + + def test_scalar(self, xp: ArrayNamespace): + a = xp.asarray(1.0) + assert nanmax(a) == 1.0 + + @pytest.mark.filterwarnings("ignore:.*All-NaN slice*.:RuntimeWarning") + def test_all_nan_slice_2d(self, xp: ArrayNamespace): + a = xp.asarray( + [ + [xp.nan, 5.0], + [xp.nan, 2.0], + ] + ) + + res = nanmax(a, axis=0, xp=xp) + expected = xp.asarray([xp.nan, 5.0]) + assert_equal(res, expected) + + @pytest.mark.skip_xp_backend( + Backend.TORCH, reason="torch.nanmax does not support tensors on meta device" + ) + @pytest.mark.parametrize("axis", [None, 0, 1]) + def test_device(self, axis: int | None, xp: ArrayNamespace, device: Device): + a = xp.asarray([[4, xp.nan, 1], [2, 5, xp.nan]], device=device) + res = nanmax(a, axis=axis) + assert get_device(res) == device + + @pytest.mark.parametrize( + ("axis", "expected_list"), [(0, [4.0, 3.0, 1.0]), (1, [4.0, 3.0])] + ) + def test_xp(self, axis: int | None, expected_list: list[float], xp: ArrayNamespace): + a = xp.asarray([[4, xp.nan, 1], [2, 3, xp.nan]]) + res = nanmax(a, axis=axis, xp=xp) + expected = xp.asarray(expected_list) + assert_equal(res, expected) + + +class TestNanSum: + def test_simple(self, xp: ArrayNamespace): + a = xp.asarray([[1.0, 2.0], [3.0, xp.nan]]) + + res = nansum(a) + expected = 6.0 + assert res == expected + + res = nansum(a, axis=0) + expected = xp.asarray([4.0, 2.0]) + assert_equal(res, expected) + + res = nansum(a, axis=1) + expected = xp.asarray([3.0, 3.0]) + assert_equal(res, expected) + + def test_bigger(self, xp: ArrayNamespace): + a = xp.asarray( + [ + [1.0, xp.nan, 4.0, 5.0], + [xp.nan, -2.0, xp.nan, -4.0], + [2.0, 1.0, 3.0, xp.nan], + ] + ) + + res = nansum(a, axis=0) + expected = xp.asarray([3.0, -1.0, 7.0, 1.0]) + assert_equal(res, expected) + + res = nansum(a, axis=1) + expected = xp.asarray([10.0, -6.0, 6.0]) + assert_equal(res, expected) + + def test_all_nan_slice(self, xp: ArrayNamespace): + a = xp.asarray([[xp.nan, 1.0], [xp.nan, xp.nan]]) + + res = nansum(a, axis=0) + expected = xp.asarray([0.0, 1.0]) + assert_equal(res, expected) + + def test_scalar(self, xp: ArrayNamespace): + a = xp.asarray(1.0) + assert nansum(a) == 1.0 + + @pytest.mark.skip_xp_backend( + Backend.TORCH, reason="torch.nansum does not support tensors on meta device" + ) + @pytest.mark.parametrize("axis", [None, 0, 1]) + def test_device(self, axis: int | None, xp: ArrayNamespace, device: Device): + a = xp.asarray([[4.0, xp.nan, 1.0], [2.0, 5.0, xp.nan]], device=device) + res = nansum(a, axis=axis) + assert get_device(res) == device + + @pytest.mark.parametrize( + ("axis", "expected_list"), [(0, [6.0, 3.0, 1.0]), (1, [5.0, 5.0])] + ) + def test_xp(self, axis: int | None, expected_list: list[float], xp: ArrayNamespace): + a = xp.asarray([[4.0, xp.nan, 1.0], [2.0, 3.0, xp.nan]]) + res = nansum(a, axis=axis, xp=xp) + expected = xp.asarray(expected_list) + assert_equal(res, expected) diff --git a/tests/test_testing.py b/tests/test_testing.py index 5be1c8fe..f2e01af3 100644 --- a/tests/test_testing.py +++ b/tests/test_testing.py @@ -8,14 +8,13 @@ from typing_extensions import override from array_api_extra._lib._backends import Backend -from array_api_extra._lib._utils._compat import ( +from array_api_extra._lib._compat import ( array_namespace, is_dask_namespace, is_jax_namespace, ) -from array_api_extra._lib._utils._typing import Array, ArrayNamespace, Device +from array_api_extra._lib._typing import Array, ArrayNamespace, Device from array_api_extra.testing import ( - _as_numpy_array, assert_close, assert_close_nulp, assert_equal, @@ -23,6 +22,7 @@ lazy_xp_function, patch_lazy_xp_functions, ) +from array_api_extra.testing._testing import _as_numpy_array # pyright: reportUnknownParameterType=false,reportMissingParameterType=false diff --git a/vendor_tests/test_vendor.py b/vendor_tests/test_vendor.py index bac073a2..51ac821e 100644 --- a/vendor_tests/test_vendor.py +++ b/vendor_tests/test_vendor.py @@ -79,7 +79,7 @@ def f(x: Any) -> Any: def test_vendor_extra_uses_vendor_compat(): from ._array_api_compat_vendor import array_namespace as n1 - from .array_api_extra._lib._utils._compat import ( # type: ignore[import-not-found] + from .array_api_extra._lib._compat import ( # type: ignore[import-not-found] array_namespace as n2, )