"""Delegation to existing implementations for Public API Functions."""
from collections.abc import Sequence
from types import ModuleType
from typing import Literal
from ._lib import _funcs
from ._lib._utils._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._utils._compat import device as get_device
from ._lib._utils._helpers import asarrays, eager_shape
from ._lib._utils._typing import Array, DType
__all__ = [
"atleast_nd",
"cov",
"create_diagonal",
"expand_dims",
"isclose",
"nan_to_num",
"one_hot",
"pad",
"sinc",
]
def atleast_nd(x: Array, /, *, ndim: int, xp: ModuleType | 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)
def cov(m: Array, /, *, xp: ModuleType | 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`.
With the exception of supporting batch input, this provides a subset of
the functionality of ``numpy.cov``.
Parameters
----------
m : array
An array of shape ``(..., N, M)`` whose innermost two dimensions
contain *M* observations of *N* variables. That is,
each row of `m` represents a variable, and each column a single
observation of all those variables.
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.
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)
"""
if xp is None:
xp = array_namespace(m)
if (
is_numpy_namespace(xp)
or is_cupy_namespace(xp)
or is_torch_namespace(xp)
or is_dask_namespace(xp)
or is_jax_namespace(xp)
) and m.ndim <= 2:
return xp.cov(m)
return _funcs.cov(m, xp=xp)
def create_diagonal(
x: Array, /, *, offset: int = 0, xp: ModuleType | 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)) and x.ndim < 2:
return xp.diag(x, k=offset)
if (is_jax_namespace(xp) or is_numpy_namespace(xp)) and x.ndim < 3:
batch_dim, n = eager_shape(x)[:-1], eager_shape(x, -1)[0] + abs(offset)
return xp.reshape(xp.diag(x, k=offset), (*batch_dim, n, n))
return _funcs.create_diagonal(x, offset=offset, xp=xp)
def expand_dims(
a: Array, /, *, axis: int | tuple[int, ...] = (0,), xp: ModuleType | None = None
) -> Array:
"""
Expand the shape of an array.
Insert (a) new axis/axes that will appear at the position(s) specified by
`axis` in the expanded array shape.
This is ``xp.expand_dims`` for `axis` an int *or a tuple of ints*.
Roughly equivalent to ``numpy.expand_dims`` for NumPy arrays.
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: ModuleType | 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 nan_to_num(
x: Array | float | complex,
/,
*,
fill_value: int | float = 0.0,
xp: ModuleType | 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 one_hot(
x: Array,
/,
num_classes: int,
*,
dtype: DType | None = None,
axis: int = -1,
xp: ModuleType | 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: ModuleType | 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)
# https://github.com/pytorch/pytorch/blob/cf76c05b4dc629ac989d1fb8e789d4fac04a095a/torch/_numpy/_funcs_impl.py#L2045-L2056
if is_torch_namespace(xp):
pad_width = xp.asarray(pad_width)
pad_width = xp.broadcast_to(pad_width, (x.ndim, 2))
pad_width = xp.flip(pad_width, axis=(0,)).flatten()
return xp.nn.functional.pad(x, tuple(pad_width), value=constant_values) # type: ignore[arg-type] # pyright: ignore[reportArgumentType]
return _funcs.pad(x, pad_width, constant_values=constant_values, xp=xp)
def setdiff1d(
x1: Array | complex,
x2: Array | complex,
/,
*,
assume_unique: bool = False,
xp: ModuleType | 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: ModuleType | 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: ModuleType | 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: ModuleType | 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: ModuleType | None = None,
) -> Array:
"""
Determine whether each element in `a` is present in `b`.
Return a boolean array of the same shape as `a` that is True for elements
that are in `b` and False otherwise.
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)