1134 lines
43 KiB
Python
1134 lines
43 KiB
Python
"""
|
|
Wrapper class around the ndarray object for the array API standard.
|
|
|
|
The array API standard defines some behaviors differently than ndarray, in
|
|
particular, type promotion rules are different (the standard has no
|
|
value-based casting). The standard also specifies a more limited subset of
|
|
array methods and functionalities than are implemented on ndarray. Since the
|
|
goal of the array_api namespace is to be a minimal implementation of the array
|
|
API standard, we need to define a separate wrapper class for the array_api
|
|
namespace.
|
|
|
|
The standard compliant class is only a wrapper class. It is *not* a subclass
|
|
of ndarray.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import operator
|
|
from enum import IntEnum
|
|
from ._creation_functions import asarray
|
|
from ._dtypes import (
|
|
_all_dtypes,
|
|
_boolean_dtypes,
|
|
_integer_dtypes,
|
|
_integer_or_boolean_dtypes,
|
|
_floating_dtypes,
|
|
_complex_floating_dtypes,
|
|
_numeric_dtypes,
|
|
_result_type,
|
|
_dtype_categories,
|
|
)
|
|
|
|
from typing import TYPE_CHECKING, Optional, Tuple, Union, Any, SupportsIndex
|
|
import types
|
|
|
|
if TYPE_CHECKING:
|
|
from ._typing import Any, PyCapsule, Device, Dtype
|
|
import numpy.typing as npt
|
|
|
|
import numpy as np
|
|
|
|
from numpy import array_api
|
|
|
|
|
|
class Array:
|
|
"""
|
|
n-d array object for the array API namespace.
|
|
|
|
See the docstring of :py:obj:`np.ndarray <numpy.ndarray>` for more
|
|
information.
|
|
|
|
This is a wrapper around numpy.ndarray that restricts the usage to only
|
|
those things that are required by the array API namespace. Note,
|
|
attributes on this object that start with a single underscore are not part
|
|
of the API specification and should only be used internally. This object
|
|
should not be constructed directly. Rather, use one of the creation
|
|
functions, such as asarray().
|
|
|
|
"""
|
|
_array: np.ndarray[Any, Any]
|
|
|
|
# Use a custom constructor instead of __init__, as manually initializing
|
|
# this class is not supported API.
|
|
@classmethod
|
|
def _new(cls, x, /):
|
|
"""
|
|
This is a private method for initializing the array API Array
|
|
object.
|
|
|
|
Functions outside of the array_api submodule should not use this
|
|
method. Use one of the creation functions instead, such as
|
|
``asarray``.
|
|
|
|
"""
|
|
obj = super().__new__(cls)
|
|
# Note: The spec does not have array scalars, only 0-D arrays.
|
|
if isinstance(x, np.generic):
|
|
# Convert the array scalar to a 0-D array
|
|
x = np.asarray(x)
|
|
if x.dtype not in _all_dtypes:
|
|
raise TypeError(
|
|
f"The array_api namespace does not support the dtype '{x.dtype}'"
|
|
)
|
|
obj._array = x
|
|
return obj
|
|
|
|
# Prevent Array() from working
|
|
def __new__(cls, *args, **kwargs):
|
|
raise TypeError(
|
|
"The array_api Array object should not be instantiated directly. Use an array creation function, such as asarray(), instead."
|
|
)
|
|
|
|
# These functions are not required by the spec, but are implemented for
|
|
# the sake of usability.
|
|
|
|
def __str__(self: Array, /) -> str:
|
|
"""
|
|
Performs the operation __str__.
|
|
"""
|
|
return self._array.__str__().replace("array", "Array")
|
|
|
|
def __repr__(self: Array, /) -> str:
|
|
"""
|
|
Performs the operation __repr__.
|
|
"""
|
|
suffix = f", dtype={self.dtype.name})"
|
|
if 0 in self.shape:
|
|
prefix = "empty("
|
|
mid = str(self.shape)
|
|
else:
|
|
prefix = "Array("
|
|
mid = np.array2string(self._array, separator=', ', prefix=prefix, suffix=suffix)
|
|
return prefix + mid + suffix
|
|
|
|
# This function is not required by the spec, but we implement it here for
|
|
# convenience so that np.asarray(np.array_api.Array) will work.
|
|
def __array__(self, dtype: None | np.dtype[Any] = None) -> npt.NDArray[Any]:
|
|
"""
|
|
Warning: this method is NOT part of the array API spec. Implementers
|
|
of other libraries need not include it, and users should not assume it
|
|
will be present in other implementations.
|
|
|
|
"""
|
|
return np.asarray(self._array, dtype=dtype)
|
|
|
|
# These are various helper functions to make the array behavior match the
|
|
# spec in places where it either deviates from or is more strict than
|
|
# NumPy behavior
|
|
|
|
def _check_allowed_dtypes(self, other: bool | int | float | Array, dtype_category: str, op: str) -> Array:
|
|
"""
|
|
Helper function for operators to only allow specific input dtypes
|
|
|
|
Use like
|
|
|
|
other = self._check_allowed_dtypes(other, 'numeric', '__add__')
|
|
if other is NotImplemented:
|
|
return other
|
|
"""
|
|
|
|
if self.dtype not in _dtype_categories[dtype_category]:
|
|
raise TypeError(f"Only {dtype_category} dtypes are allowed in {op}")
|
|
if isinstance(other, (int, complex, float, bool)):
|
|
other = self._promote_scalar(other)
|
|
elif isinstance(other, Array):
|
|
if other.dtype not in _dtype_categories[dtype_category]:
|
|
raise TypeError(f"Only {dtype_category} dtypes are allowed in {op}")
|
|
else:
|
|
return NotImplemented
|
|
|
|
# This will raise TypeError for type combinations that are not allowed
|
|
# to promote in the spec (even if the NumPy array operator would
|
|
# promote them).
|
|
res_dtype = _result_type(self.dtype, other.dtype)
|
|
if op.startswith("__i"):
|
|
# Note: NumPy will allow in-place operators in some cases where
|
|
# the type promoted operator does not match the left-hand side
|
|
# operand. For example,
|
|
|
|
# >>> a = np.array(1, dtype=np.int8)
|
|
# >>> a += np.array(1, dtype=np.int16)
|
|
|
|
# The spec explicitly disallows this.
|
|
if res_dtype != self.dtype:
|
|
raise TypeError(
|
|
f"Cannot perform {op} with dtypes {self.dtype} and {other.dtype}"
|
|
)
|
|
|
|
return other
|
|
|
|
# Helper function to match the type promotion rules in the spec
|
|
def _promote_scalar(self, scalar):
|
|
"""
|
|
Returns a promoted version of a Python scalar appropriate for use with
|
|
operations on self.
|
|
|
|
This may raise an OverflowError in cases where the scalar is an
|
|
integer that is too large to fit in a NumPy integer dtype, or
|
|
TypeError when the scalar type is incompatible with the dtype of self.
|
|
"""
|
|
# Note: Only Python scalar types that match the array dtype are
|
|
# allowed.
|
|
if isinstance(scalar, bool):
|
|
if self.dtype not in _boolean_dtypes:
|
|
raise TypeError(
|
|
"Python bool scalars can only be promoted with bool arrays"
|
|
)
|
|
elif isinstance(scalar, int):
|
|
if self.dtype in _boolean_dtypes:
|
|
raise TypeError(
|
|
"Python int scalars cannot be promoted with bool arrays"
|
|
)
|
|
if self.dtype in _integer_dtypes:
|
|
info = np.iinfo(self.dtype)
|
|
if not (info.min <= scalar <= info.max):
|
|
raise OverflowError(
|
|
"Python int scalars must be within the bounds of the dtype for integer arrays"
|
|
)
|
|
# int + array(floating) is allowed
|
|
elif isinstance(scalar, float):
|
|
if self.dtype not in _floating_dtypes:
|
|
raise TypeError(
|
|
"Python float scalars can only be promoted with floating-point arrays."
|
|
)
|
|
elif isinstance(scalar, complex):
|
|
if self.dtype not in _complex_floating_dtypes:
|
|
raise TypeError(
|
|
"Python complex scalars can only be promoted with complex floating-point arrays."
|
|
)
|
|
else:
|
|
raise TypeError("'scalar' must be a Python scalar")
|
|
|
|
# Note: scalars are unconditionally cast to the same dtype as the
|
|
# array.
|
|
|
|
# Note: the spec only specifies integer-dtype/int promotion
|
|
# behavior for integers within the bounds of the integer dtype.
|
|
# Outside of those bounds we use the default NumPy behavior (either
|
|
# cast or raise OverflowError).
|
|
return Array._new(np.array(scalar, self.dtype))
|
|
|
|
@staticmethod
|
|
def _normalize_two_args(x1, x2) -> Tuple[Array, Array]:
|
|
"""
|
|
Normalize inputs to two arg functions to fix type promotion rules
|
|
|
|
NumPy deviates from the spec type promotion rules in cases where one
|
|
argument is 0-dimensional and the other is not. For example:
|
|
|
|
>>> import numpy as np
|
|
>>> a = np.array([1.0], dtype=np.float32)
|
|
>>> b = np.array(1.0, dtype=np.float64)
|
|
>>> np.add(a, b) # The spec says this should be float64
|
|
array([2.], dtype=float32)
|
|
|
|
To fix this, we add a dimension to the 0-dimension array before passing it
|
|
through. This works because a dimension would be added anyway from
|
|
broadcasting, so the resulting shape is the same, but this prevents NumPy
|
|
from not promoting the dtype.
|
|
"""
|
|
# Another option would be to use signature=(x1.dtype, x2.dtype, None),
|
|
# but that only works for ufuncs, so we would have to call the ufuncs
|
|
# directly in the operator methods. One should also note that this
|
|
# sort of trick wouldn't work for functions like searchsorted, which
|
|
# don't do normal broadcasting, but there aren't any functions like
|
|
# that in the array API namespace.
|
|
if x1.ndim == 0 and x2.ndim != 0:
|
|
# The _array[None] workaround was chosen because it is relatively
|
|
# performant. broadcast_to(x1._array, x2.shape) is much slower. We
|
|
# could also manually type promote x2, but that is more complicated
|
|
# and about the same performance as this.
|
|
x1 = Array._new(x1._array[None])
|
|
elif x2.ndim == 0 and x1.ndim != 0:
|
|
x2 = Array._new(x2._array[None])
|
|
return (x1, x2)
|
|
|
|
# Note: A large fraction of allowed indices are disallowed here (see the
|
|
# docstring below)
|
|
def _validate_index(self, key):
|
|
"""
|
|
Validate an index according to the array API.
|
|
|
|
The array API specification only requires a subset of indices that are
|
|
supported by NumPy. This function will reject any index that is
|
|
allowed by NumPy but not required by the array API specification. We
|
|
always raise ``IndexError`` on such indices (the spec does not require
|
|
any specific behavior on them, but this makes the NumPy array API
|
|
namespace a minimal implementation of the spec). See
|
|
https://data-apis.org/array-api/latest/API_specification/indexing.html
|
|
for the full list of required indexing behavior
|
|
|
|
This function raises IndexError if the index ``key`` is invalid. It
|
|
only raises ``IndexError`` on indices that are not already rejected by
|
|
NumPy, as NumPy will already raise the appropriate error on such
|
|
indices. ``shape`` may be None, in which case, only cases that are
|
|
independent of the array shape are checked.
|
|
|
|
The following cases are allowed by NumPy, but not specified by the array
|
|
API specification:
|
|
|
|
- Indices to not include an implicit ellipsis at the end. That is,
|
|
every axis of an array must be explicitly indexed or an ellipsis
|
|
included. This behaviour is sometimes referred to as flat indexing.
|
|
|
|
- The start and stop of a slice may not be out of bounds. In
|
|
particular, for a slice ``i:j:k`` on an axis of size ``n``, only the
|
|
following are allowed:
|
|
|
|
- ``i`` or ``j`` omitted (``None``).
|
|
- ``-n <= i <= max(0, n - 1)``.
|
|
- For ``k > 0`` or ``k`` omitted (``None``), ``-n <= j <= n``.
|
|
- For ``k < 0``, ``-n - 1 <= j <= max(0, n - 1)``.
|
|
|
|
- Boolean array indices are not allowed as part of a larger tuple
|
|
index.
|
|
|
|
- Integer array indices are not allowed (with the exception of 0-D
|
|
arrays, which are treated the same as scalars).
|
|
|
|
Additionally, it should be noted that indices that would return a
|
|
scalar in NumPy will return a 0-D array. Array scalars are not allowed
|
|
in the specification, only 0-D arrays. This is done in the
|
|
``Array._new`` constructor, not this function.
|
|
|
|
"""
|
|
_key = key if isinstance(key, tuple) else (key,)
|
|
for i in _key:
|
|
if isinstance(i, bool) or not (
|
|
isinstance(i, SupportsIndex) # i.e. ints
|
|
or isinstance(i, slice)
|
|
or i == Ellipsis
|
|
or i is None
|
|
or isinstance(i, Array)
|
|
or isinstance(i, np.ndarray)
|
|
):
|
|
raise IndexError(
|
|
f"Single-axes index {i} has {type(i)=}, but only "
|
|
"integers, slices (:), ellipsis (...), newaxis (None), "
|
|
"zero-dimensional integer arrays and boolean arrays "
|
|
"are specified in the Array API."
|
|
)
|
|
|
|
nonexpanding_key = []
|
|
single_axes = []
|
|
n_ellipsis = 0
|
|
key_has_mask = False
|
|
for i in _key:
|
|
if i is not None:
|
|
nonexpanding_key.append(i)
|
|
if isinstance(i, Array) or isinstance(i, np.ndarray):
|
|
if i.dtype in _boolean_dtypes:
|
|
key_has_mask = True
|
|
single_axes.append(i)
|
|
else:
|
|
# i must not be an array here, to avoid elementwise equals
|
|
if i == Ellipsis:
|
|
n_ellipsis += 1
|
|
else:
|
|
single_axes.append(i)
|
|
|
|
n_single_axes = len(single_axes)
|
|
if n_ellipsis > 1:
|
|
return # handled by ndarray
|
|
elif n_ellipsis == 0:
|
|
# Note boolean masks must be the sole index, which we check for
|
|
# later on.
|
|
if not key_has_mask and n_single_axes < self.ndim:
|
|
raise IndexError(
|
|
f"{self.ndim=}, but the multi-axes index only specifies "
|
|
f"{n_single_axes} dimensions. If this was intentional, "
|
|
"add a trailing ellipsis (...) which expands into as many "
|
|
"slices (:) as necessary - this is what np.ndarray arrays "
|
|
"implicitly do, but such flat indexing behaviour is not "
|
|
"specified in the Array API."
|
|
)
|
|
|
|
if n_ellipsis == 0:
|
|
indexed_shape = self.shape
|
|
else:
|
|
ellipsis_start = None
|
|
for pos, i in enumerate(nonexpanding_key):
|
|
if not (isinstance(i, Array) or isinstance(i, np.ndarray)):
|
|
if i == Ellipsis:
|
|
ellipsis_start = pos
|
|
break
|
|
assert ellipsis_start is not None # sanity check
|
|
ellipsis_end = self.ndim - (n_single_axes - ellipsis_start)
|
|
indexed_shape = (
|
|
self.shape[:ellipsis_start] + self.shape[ellipsis_end:]
|
|
)
|
|
for i, side in zip(single_axes, indexed_shape):
|
|
if isinstance(i, slice):
|
|
if side == 0:
|
|
f_range = "0 (or None)"
|
|
else:
|
|
f_range = f"between -{side} and {side - 1} (or None)"
|
|
if i.start is not None:
|
|
try:
|
|
start = operator.index(i.start)
|
|
except TypeError:
|
|
pass # handled by ndarray
|
|
else:
|
|
if not (-side <= start <= side):
|
|
raise IndexError(
|
|
f"Slice {i} contains {start=}, but should be "
|
|
f"{f_range} for an axis of size {side} "
|
|
"(out-of-bounds starts are not specified in "
|
|
"the Array API)"
|
|
)
|
|
if i.stop is not None:
|
|
try:
|
|
stop = operator.index(i.stop)
|
|
except TypeError:
|
|
pass # handled by ndarray
|
|
else:
|
|
if not (-side <= stop <= side):
|
|
raise IndexError(
|
|
f"Slice {i} contains {stop=}, but should be "
|
|
f"{f_range} for an axis of size {side} "
|
|
"(out-of-bounds stops are not specified in "
|
|
"the Array API)"
|
|
)
|
|
elif isinstance(i, Array):
|
|
if i.dtype in _boolean_dtypes and len(_key) != 1:
|
|
assert isinstance(key, tuple) # sanity check
|
|
raise IndexError(
|
|
f"Single-axes index {i} is a boolean array and "
|
|
f"{len(key)=}, but masking is only specified in the "
|
|
"Array API when the array is the sole index."
|
|
)
|
|
elif i.dtype in _integer_dtypes and i.ndim != 0:
|
|
raise IndexError(
|
|
f"Single-axes index {i} is a non-zero-dimensional "
|
|
"integer array, but advanced integer indexing is not "
|
|
"specified in the Array API."
|
|
)
|
|
elif isinstance(i, tuple):
|
|
raise IndexError(
|
|
f"Single-axes index {i} is a tuple, but nested tuple "
|
|
"indices are not specified in the Array API."
|
|
)
|
|
|
|
# Everything below this line is required by the spec.
|
|
|
|
def __abs__(self: Array, /) -> Array:
|
|
"""
|
|
Performs the operation __abs__.
|
|
"""
|
|
if self.dtype not in _numeric_dtypes:
|
|
raise TypeError("Only numeric dtypes are allowed in __abs__")
|
|
res = self._array.__abs__()
|
|
return self.__class__._new(res)
|
|
|
|
def __add__(self: Array, other: Union[int, float, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __add__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "numeric", "__add__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self, other = self._normalize_two_args(self, other)
|
|
res = self._array.__add__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __and__(self: Array, other: Union[int, bool, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __and__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "integer or boolean", "__and__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self, other = self._normalize_two_args(self, other)
|
|
res = self._array.__and__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __array_namespace__(
|
|
self: Array, /, *, api_version: Optional[str] = None
|
|
) -> types.ModuleType:
|
|
if api_version is not None and not api_version.startswith("2021."):
|
|
raise ValueError(f"Unrecognized array API version: {api_version!r}")
|
|
return array_api
|
|
|
|
def __bool__(self: Array, /) -> bool:
|
|
"""
|
|
Performs the operation __bool__.
|
|
"""
|
|
# Note: This is an error here.
|
|
if self._array.ndim != 0:
|
|
raise TypeError("bool is only allowed on arrays with 0 dimensions")
|
|
res = self._array.__bool__()
|
|
return res
|
|
|
|
def __complex__(self: Array, /) -> complex:
|
|
"""
|
|
Performs the operation __complex__.
|
|
"""
|
|
# Note: This is an error here.
|
|
if self._array.ndim != 0:
|
|
raise TypeError("complex is only allowed on arrays with 0 dimensions")
|
|
res = self._array.__complex__()
|
|
return res
|
|
|
|
def __dlpack__(self: Array, /, *, stream: None = None) -> PyCapsule:
|
|
"""
|
|
Performs the operation __dlpack__.
|
|
"""
|
|
return self._array.__dlpack__(stream=stream)
|
|
|
|
def __dlpack_device__(self: Array, /) -> Tuple[IntEnum, int]:
|
|
"""
|
|
Performs the operation __dlpack_device__.
|
|
"""
|
|
# Note: device support is required for this
|
|
return self._array.__dlpack_device__()
|
|
|
|
def __eq__(self: Array, other: Union[int, float, bool, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __eq__.
|
|
"""
|
|
# Even though "all" dtypes are allowed, we still require them to be
|
|
# promotable with each other.
|
|
other = self._check_allowed_dtypes(other, "all", "__eq__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self, other = self._normalize_two_args(self, other)
|
|
res = self._array.__eq__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __float__(self: Array, /) -> float:
|
|
"""
|
|
Performs the operation __float__.
|
|
"""
|
|
# Note: This is an error here.
|
|
if self._array.ndim != 0:
|
|
raise TypeError("float is only allowed on arrays with 0 dimensions")
|
|
if self.dtype in _complex_floating_dtypes:
|
|
raise TypeError("float is not allowed on complex floating-point arrays")
|
|
res = self._array.__float__()
|
|
return res
|
|
|
|
def __floordiv__(self: Array, other: Union[int, float, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __floordiv__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "real numeric", "__floordiv__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self, other = self._normalize_two_args(self, other)
|
|
res = self._array.__floordiv__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __ge__(self: Array, other: Union[int, float, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __ge__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "real numeric", "__ge__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self, other = self._normalize_two_args(self, other)
|
|
res = self._array.__ge__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __getitem__(
|
|
self: Array,
|
|
key: Union[
|
|
int,
|
|
slice,
|
|
ellipsis,
|
|
Tuple[Union[int, slice, ellipsis, None], ...],
|
|
Array,
|
|
],
|
|
/,
|
|
) -> Array:
|
|
"""
|
|
Performs the operation __getitem__.
|
|
"""
|
|
# Note: Only indices required by the spec are allowed. See the
|
|
# docstring of _validate_index
|
|
self._validate_index(key)
|
|
if isinstance(key, Array):
|
|
# Indexing self._array with array_api arrays can be erroneous
|
|
key = key._array
|
|
res = self._array.__getitem__(key)
|
|
return self._new(res)
|
|
|
|
def __gt__(self: Array, other: Union[int, float, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __gt__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "real numeric", "__gt__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self, other = self._normalize_two_args(self, other)
|
|
res = self._array.__gt__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __int__(self: Array, /) -> int:
|
|
"""
|
|
Performs the operation __int__.
|
|
"""
|
|
# Note: This is an error here.
|
|
if self._array.ndim != 0:
|
|
raise TypeError("int is only allowed on arrays with 0 dimensions")
|
|
if self.dtype in _complex_floating_dtypes:
|
|
raise TypeError("int is not allowed on complex floating-point arrays")
|
|
res = self._array.__int__()
|
|
return res
|
|
|
|
def __index__(self: Array, /) -> int:
|
|
"""
|
|
Performs the operation __index__.
|
|
"""
|
|
res = self._array.__index__()
|
|
return res
|
|
|
|
def __invert__(self: Array, /) -> Array:
|
|
"""
|
|
Performs the operation __invert__.
|
|
"""
|
|
if self.dtype not in _integer_or_boolean_dtypes:
|
|
raise TypeError("Only integer or boolean dtypes are allowed in __invert__")
|
|
res = self._array.__invert__()
|
|
return self.__class__._new(res)
|
|
|
|
def __le__(self: Array, other: Union[int, float, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __le__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "real numeric", "__le__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self, other = self._normalize_two_args(self, other)
|
|
res = self._array.__le__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __lshift__(self: Array, other: Union[int, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __lshift__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "integer", "__lshift__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self, other = self._normalize_two_args(self, other)
|
|
res = self._array.__lshift__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __lt__(self: Array, other: Union[int, float, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __lt__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "real numeric", "__lt__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self, other = self._normalize_two_args(self, other)
|
|
res = self._array.__lt__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __matmul__(self: Array, other: Array, /) -> Array:
|
|
"""
|
|
Performs the operation __matmul__.
|
|
"""
|
|
# matmul is not defined for scalars, but without this, we may get
|
|
# the wrong error message from asarray.
|
|
other = self._check_allowed_dtypes(other, "numeric", "__matmul__")
|
|
if other is NotImplemented:
|
|
return other
|
|
res = self._array.__matmul__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __mod__(self: Array, other: Union[int, float, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __mod__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "real numeric", "__mod__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self, other = self._normalize_two_args(self, other)
|
|
res = self._array.__mod__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __mul__(self: Array, other: Union[int, float, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __mul__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "numeric", "__mul__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self, other = self._normalize_two_args(self, other)
|
|
res = self._array.__mul__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __ne__(self: Array, other: Union[int, float, bool, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __ne__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "all", "__ne__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self, other = self._normalize_two_args(self, other)
|
|
res = self._array.__ne__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __neg__(self: Array, /) -> Array:
|
|
"""
|
|
Performs the operation __neg__.
|
|
"""
|
|
if self.dtype not in _numeric_dtypes:
|
|
raise TypeError("Only numeric dtypes are allowed in __neg__")
|
|
res = self._array.__neg__()
|
|
return self.__class__._new(res)
|
|
|
|
def __or__(self: Array, other: Union[int, bool, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __or__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "integer or boolean", "__or__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self, other = self._normalize_two_args(self, other)
|
|
res = self._array.__or__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __pos__(self: Array, /) -> Array:
|
|
"""
|
|
Performs the operation __pos__.
|
|
"""
|
|
if self.dtype not in _numeric_dtypes:
|
|
raise TypeError("Only numeric dtypes are allowed in __pos__")
|
|
res = self._array.__pos__()
|
|
return self.__class__._new(res)
|
|
|
|
def __pow__(self: Array, other: Union[int, float, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __pow__.
|
|
"""
|
|
from ._elementwise_functions import pow
|
|
|
|
other = self._check_allowed_dtypes(other, "numeric", "__pow__")
|
|
if other is NotImplemented:
|
|
return other
|
|
# Note: NumPy's __pow__ does not follow type promotion rules for 0-d
|
|
# arrays, so we use pow() here instead.
|
|
return pow(self, other)
|
|
|
|
def __rshift__(self: Array, other: Union[int, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __rshift__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "integer", "__rshift__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self, other = self._normalize_two_args(self, other)
|
|
res = self._array.__rshift__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __setitem__(
|
|
self,
|
|
key: Union[
|
|
int, slice, ellipsis, Tuple[Union[int, slice, ellipsis], ...], Array
|
|
],
|
|
value: Union[int, float, bool, Array],
|
|
/,
|
|
) -> None:
|
|
"""
|
|
Performs the operation __setitem__.
|
|
"""
|
|
# Note: Only indices required by the spec are allowed. See the
|
|
# docstring of _validate_index
|
|
self._validate_index(key)
|
|
if isinstance(key, Array):
|
|
# Indexing self._array with array_api arrays can be erroneous
|
|
key = key._array
|
|
self._array.__setitem__(key, asarray(value)._array)
|
|
|
|
def __sub__(self: Array, other: Union[int, float, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __sub__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "numeric", "__sub__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self, other = self._normalize_two_args(self, other)
|
|
res = self._array.__sub__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
# PEP 484 requires int to be a subtype of float, but __truediv__ should
|
|
# not accept int.
|
|
def __truediv__(self: Array, other: Union[float, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __truediv__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "floating-point", "__truediv__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self, other = self._normalize_two_args(self, other)
|
|
res = self._array.__truediv__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __xor__(self: Array, other: Union[int, bool, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __xor__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "integer or boolean", "__xor__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self, other = self._normalize_two_args(self, other)
|
|
res = self._array.__xor__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __iadd__(self: Array, other: Union[int, float, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __iadd__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "numeric", "__iadd__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self._array.__iadd__(other._array)
|
|
return self
|
|
|
|
def __radd__(self: Array, other: Union[int, float, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __radd__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "numeric", "__radd__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self, other = self._normalize_two_args(self, other)
|
|
res = self._array.__radd__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __iand__(self: Array, other: Union[int, bool, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __iand__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "integer or boolean", "__iand__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self._array.__iand__(other._array)
|
|
return self
|
|
|
|
def __rand__(self: Array, other: Union[int, bool, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __rand__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "integer or boolean", "__rand__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self, other = self._normalize_two_args(self, other)
|
|
res = self._array.__rand__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __ifloordiv__(self: Array, other: Union[int, float, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __ifloordiv__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "real numeric", "__ifloordiv__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self._array.__ifloordiv__(other._array)
|
|
return self
|
|
|
|
def __rfloordiv__(self: Array, other: Union[int, float, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __rfloordiv__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "real numeric", "__rfloordiv__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self, other = self._normalize_two_args(self, other)
|
|
res = self._array.__rfloordiv__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __ilshift__(self: Array, other: Union[int, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __ilshift__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "integer", "__ilshift__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self._array.__ilshift__(other._array)
|
|
return self
|
|
|
|
def __rlshift__(self: Array, other: Union[int, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __rlshift__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "integer", "__rlshift__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self, other = self._normalize_two_args(self, other)
|
|
res = self._array.__rlshift__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __imatmul__(self: Array, other: Array, /) -> Array:
|
|
"""
|
|
Performs the operation __imatmul__.
|
|
"""
|
|
# matmul is not defined for scalars, but without this, we may get
|
|
# the wrong error message from asarray.
|
|
other = self._check_allowed_dtypes(other, "numeric", "__imatmul__")
|
|
if other is NotImplemented:
|
|
return other
|
|
res = self._array.__imatmul__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __rmatmul__(self: Array, other: Array, /) -> Array:
|
|
"""
|
|
Performs the operation __rmatmul__.
|
|
"""
|
|
# matmul is not defined for scalars, but without this, we may get
|
|
# the wrong error message from asarray.
|
|
other = self._check_allowed_dtypes(other, "numeric", "__rmatmul__")
|
|
if other is NotImplemented:
|
|
return other
|
|
res = self._array.__rmatmul__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __imod__(self: Array, other: Union[int, float, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __imod__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "real numeric", "__imod__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self._array.__imod__(other._array)
|
|
return self
|
|
|
|
def __rmod__(self: Array, other: Union[int, float, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __rmod__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "real numeric", "__rmod__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self, other = self._normalize_two_args(self, other)
|
|
res = self._array.__rmod__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __imul__(self: Array, other: Union[int, float, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __imul__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "numeric", "__imul__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self._array.__imul__(other._array)
|
|
return self
|
|
|
|
def __rmul__(self: Array, other: Union[int, float, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __rmul__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "numeric", "__rmul__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self, other = self._normalize_two_args(self, other)
|
|
res = self._array.__rmul__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __ior__(self: Array, other: Union[int, bool, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __ior__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "integer or boolean", "__ior__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self._array.__ior__(other._array)
|
|
return self
|
|
|
|
def __ror__(self: Array, other: Union[int, bool, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __ror__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "integer or boolean", "__ror__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self, other = self._normalize_two_args(self, other)
|
|
res = self._array.__ror__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __ipow__(self: Array, other: Union[int, float, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __ipow__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "numeric", "__ipow__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self._array.__ipow__(other._array)
|
|
return self
|
|
|
|
def __rpow__(self: Array, other: Union[int, float, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __rpow__.
|
|
"""
|
|
from ._elementwise_functions import pow
|
|
|
|
other = self._check_allowed_dtypes(other, "numeric", "__rpow__")
|
|
if other is NotImplemented:
|
|
return other
|
|
# Note: NumPy's __pow__ does not follow the spec type promotion rules
|
|
# for 0-d arrays, so we use pow() here instead.
|
|
return pow(other, self)
|
|
|
|
def __irshift__(self: Array, other: Union[int, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __irshift__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "integer", "__irshift__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self._array.__irshift__(other._array)
|
|
return self
|
|
|
|
def __rrshift__(self: Array, other: Union[int, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __rrshift__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "integer", "__rrshift__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self, other = self._normalize_two_args(self, other)
|
|
res = self._array.__rrshift__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __isub__(self: Array, other: Union[int, float, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __isub__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "numeric", "__isub__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self._array.__isub__(other._array)
|
|
return self
|
|
|
|
def __rsub__(self: Array, other: Union[int, float, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __rsub__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "numeric", "__rsub__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self, other = self._normalize_two_args(self, other)
|
|
res = self._array.__rsub__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __itruediv__(self: Array, other: Union[float, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __itruediv__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "floating-point", "__itruediv__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self._array.__itruediv__(other._array)
|
|
return self
|
|
|
|
def __rtruediv__(self: Array, other: Union[float, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __rtruediv__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "floating-point", "__rtruediv__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self, other = self._normalize_two_args(self, other)
|
|
res = self._array.__rtruediv__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def __ixor__(self: Array, other: Union[int, bool, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __ixor__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "integer or boolean", "__ixor__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self._array.__ixor__(other._array)
|
|
return self
|
|
|
|
def __rxor__(self: Array, other: Union[int, bool, Array], /) -> Array:
|
|
"""
|
|
Performs the operation __rxor__.
|
|
"""
|
|
other = self._check_allowed_dtypes(other, "integer or boolean", "__rxor__")
|
|
if other is NotImplemented:
|
|
return other
|
|
self, other = self._normalize_two_args(self, other)
|
|
res = self._array.__rxor__(other._array)
|
|
return self.__class__._new(res)
|
|
|
|
def to_device(self: Array, device: Device, /, stream: None = None) -> Array:
|
|
if stream is not None:
|
|
raise ValueError("The stream argument to to_device() is not supported")
|
|
if device == 'cpu':
|
|
return self
|
|
raise ValueError(f"Unsupported device {device!r}")
|
|
|
|
@property
|
|
def dtype(self) -> Dtype:
|
|
"""
|
|
Array API compatible wrapper for :py:meth:`np.ndarray.dtype <numpy.ndarray.dtype>`.
|
|
|
|
See its docstring for more information.
|
|
"""
|
|
return self._array.dtype
|
|
|
|
@property
|
|
def device(self) -> Device:
|
|
return "cpu"
|
|
|
|
# Note: mT is new in array API spec (see matrix_transpose)
|
|
@property
|
|
def mT(self) -> Array:
|
|
from .linalg import matrix_transpose
|
|
return matrix_transpose(self)
|
|
|
|
@property
|
|
def ndim(self) -> int:
|
|
"""
|
|
Array API compatible wrapper for :py:meth:`np.ndarray.ndim <numpy.ndarray.ndim>`.
|
|
|
|
See its docstring for more information.
|
|
"""
|
|
return self._array.ndim
|
|
|
|
@property
|
|
def shape(self) -> Tuple[int, ...]:
|
|
"""
|
|
Array API compatible wrapper for :py:meth:`np.ndarray.shape <numpy.ndarray.shape>`.
|
|
|
|
See its docstring for more information.
|
|
"""
|
|
return self._array.shape
|
|
|
|
@property
|
|
def size(self) -> int:
|
|
"""
|
|
Array API compatible wrapper for :py:meth:`np.ndarray.size <numpy.ndarray.size>`.
|
|
|
|
See its docstring for more information.
|
|
"""
|
|
return self._array.size
|
|
|
|
@property
|
|
def T(self) -> Array:
|
|
"""
|
|
Array API compatible wrapper for :py:meth:`np.ndarray.T <numpy.ndarray.T>`.
|
|
|
|
See its docstring for more information.
|
|
"""
|
|
# Note: T only works on 2-dimensional arrays. See the corresponding
|
|
# note in the specification:
|
|
# https://data-apis.org/array-api/latest/API_specification/array_object.html#t
|
|
if self.ndim != 2:
|
|
raise ValueError("x.T requires x to have 2 dimensions. Use x.mT to transpose stacks of matrices and permute_dims() to permute dimensions.")
|
|
return self.__class__._new(self._array.T)
|