diff --git a/README.md b/README.md index 9d888663..1f5b7c40 100644 --- a/README.md +++ b/README.md @@ -1,18 +1,28 @@ -# NumPy Array API compatibility library +# Array API compatibility library -This is a small wrapper around NumPy that is compatible with the [Array API -standard](https://data-apis.org/array-api/latest/). See also [NEP 47](https://numpy.org/neps/nep-0047-array-api-standard.html). +This is a small wrapper around NumPy and CuPy that is compatible with the +[Array API standard](https://data-apis.org/array-api/latest/). See also [NEP +47](https://numpy.org/neps/nep-0047-array-api-standard.html). Unlike `numpy.array_api`, this is not a strict minimal implementation of the -Array API, but rather just an extension of the main NumPy namespace with -changes needed to be compliant with the Array API. See -https://numpy.org/doc/stable/reference/array_api.html for a full list of +Array API, but rather just an extension of the main NumPy and CuPy namespaces +with changes needed to be compliant with the Array API. + +Library authors using the Array API may wish to test against `numpy.array_api` +to ensure they are not using functionality outside of the standard, but prefer +this implementation for the default when working with NumPy or CuPy arrays. + +See https://numpy.org/doc/stable/reference/array_api.html for a full list of changes. In particular, unlike `numpy.array_api`, this package does not use a separate Array object, but rather just uses `numpy.ndarray` directly. Note that some of the functionality in this library is backwards incompatible with NumPy. +This library also supports CuPy in addition to NumPy. If you want support for +other array libraries, please [open an +issue](https://github.com/data-apis/array-api-compat/issues). + Library authors using the Array API may wish to test against `numpy.array_api` to ensure they are not using functionality outside of the standard, but prefer this implementation for end users who use NumPy arrays. @@ -28,5 +38,145 @@ import numpy as np with ```py -import numpy_array_api_compat as np +import array_api_compat.numpy as np ``` + +and replace + +```py +import cupy as cp +``` + +with + +```py +import array_api_compat.cupy as cp +``` + +Each will include all the functions from the normal NumPy/CuPy namespace, +except that functions that are part of the array API are wrapped so that they +have the correct array API behavior. In each case, the array object used will +be thew same array object from the wrapped library. + + +## Helper Functions + +In addition to the default NumPy/CuPy namespace and functions in the array API +specification, there are several helper functions +included that aren't part of the specification but which are useful for using +the array API: + +- `is_array_api_obj(x)`: Return `True` if `x` is an array API compatible array + object. + +- `get_namespace(*xs)`: Get the corresponding array API namespace for the + arrays `xs`. If the arrays are NumPy or CuPy arrays, the returned namespace + will be `array_api_compat.numpy` or `array_api_compat.cupy` so that it is + array API compatible. + +- `device(x)`: Equivalent to + [`x.device`](https://data-apis.org/array-api/latest/API_specification/generated/signatures.array_object.array.device.html) + in the array API specification. Included because `numpy.ndarray` does not + include the `device` attribute and this library does not wrap or extend the + array object. Note that for NumPy, `device` is always `"cpu"`. + +- `to_device(x, device, /, *, stream=None)`: Equivalent to + [`x.to_device`](https://data-apis.org/array-api/latest/API_specification/generated/signatures.array_object.array.to_device.html). + Included because neither NumPy's nor CuPy's ndarray objects include this + method. For NumPy, this function effectively does nothing since the only + supported device is the CPU, but for CuPy, this method supports CuPy CUDA + [Device](https://docs.cupy.dev/en/stable/reference/generated/cupy.cuda.Device.html) + and + [Stream](https://docs.cupy.dev/en/stable/reference/generated/cupy.cuda.Stream.html) + objects. + +## Known Differences from the Array API Specification + +There are some known differences between this library and the array API +specification: + +- The array methods `__array_namespace__`, `device` (for NumPy), `to_device`, + and `mT` are not defined. This reuses `np.ndarray` and `cp.ndarray` and we + don't want to monkeypatch or wrap it. The helper functions `device()` and + `to_device()` are provided to work around these missing methods (see above). + `x.mT` can be replaced with `xp.linalg.matrix_transpose(x)`. + `get_namespace(x)` should be used instead of `x.__array_namespace__`. + +- NumPy value-based casting for scalars will be in effect unless explicitly + disabled with the environment variable NPY_PROMOTION_STATE=weak or + np._set_promotion_state('weak') (requires NumPy 1.24 or newer, see NEP 50 + and https://github.com/numpy/numpy/issues/22341) + +- Functions which are not wrapped may not have the same type annotations + as the spec. + +- Functions which are not wrapped may not use positional-only arguments. + +## Vendoring + +This library supports vendoring as an installation method. To vendor the +library, simply copy `array_api_compat` into the appropriate place in the +library, like + +``` +cp -R array_api_compat/ mylib/vendored/array_api_compat +``` + +You may also rename it to something else if you like (nowhere in the code +references the name "array_api_compat"). + +Alternatively, the library may be installed as dependency on PyPI. + +## Implementation + +As noted before, the goal of this library is to reuse the NumPy and CuPy array +objects, rather than wrapping or extending them. This means that the functions +need to accept and return `np.ndarray` for NumPy and `cp.ndarray` for CuPy. + +Each namespace (`array_api_compat.numpy` and `array_api_compat.cupy`) is +populated with the normal library namespace (like `from numpy import *`). Then +specific functions are replaced with wrapped variants. Wrapped functions that +have the same logic between NumPy and CuPy (which is most functions) are in +`array_api_compat/common/`. These functions are defined like + +```py +# In array_api_compat/common/_aliases.py + +def acos(x, /, xp): + return xp.arccos(x) +``` + +The `xp` argument refers to the original array namespace (either `numpy` or +`cupy`). Then in the specific `array_api_compat/numpy` and +`array_api_compat/cupy` namespace, the `get_xp` decorator is applied to these +functions, which automatically removes the `xp` argument from the function +signature and replaces it with the corresponding array library, like + +```py +# In array_api_compat/numpy/_aliases.py + +from ..common import _aliases + +import numpy as np + +acos = get_xp(np)(_aliases.acos) +``` + +This `acos` now has the signature `acos(x, /)` and calls `numpy.arccos`. + +Similarly, for CuPy: + +```py +# In array_api_compat/cupy/_aliases.py + +from ..common import _aliases + +import cupy as cp + +acos = get_xp(cp)(_aliases.acos) +``` + +Since NumPy and CuPy are nearly identical in their behaviors, this allows +writing the wrapping logic for both libraries only once. If support is added +for other libraries which differ significantly from NumPy, their wrapper code +should go in their specific sub-namespace instead of `common/`. diff --git a/array_api_compat/__init__.py b/array_api_compat/__init__.py new file mode 100644 index 00000000..34bd7565 --- /dev/null +++ b/array_api_compat/__init__.py @@ -0,0 +1,20 @@ +""" +NumPy Array API compatibility library + +This is a small wrapper around NumPy and CuPy that is compatible with the +Array API standard https://data-apis.org/array-api/latest/. See also NEP 47 +https://numpy.org/neps/nep-0047-array-api-standard.html. + +Unlike numpy.array_api, this is not a strict minimal implementation of the +Array API, but rather just an extension of the main NumPy namespace with +changes needed to be compliant with the Array API. See +https://numpy.org/doc/stable/reference/array_api.html for a full list of +changes. In particular, unlike numpy.array_api, this package does not use a +separate Array object, but rather just uses numpy.ndarray directly. + +Library authors using the Array API may wish to test against numpy.array_api +to ensure they are not using functionality outside of the standard, but prefer +this implementation for the default when working with NumPy arrays. + +""" +from .common import * diff --git a/array_api_compat/_internal.py b/array_api_compat/_internal.py new file mode 100644 index 00000000..553c0356 --- /dev/null +++ b/array_api_compat/_internal.py @@ -0,0 +1,43 @@ +""" +Internal helpers +""" + +from functools import wraps +from inspect import signature + +def get_xp(xp): + """ + Decorator to automatically replace xp with the corresponding array module. + + Use like + + import numpy as np + + @get_xp(np) + def func(x, /, xp, kwarg=None): + return xp.func(x, kwarg=kwarg) + + Note that xp must be a keyword argument and come after all non-keyword + arguments. + + """ + def inner(f): + @wraps(f) + def wrapped_f(*args, **kwargs): + return f(*args, xp=xp, **kwargs) + + sig = signature(f) + new_sig = sig.replace(parameters=[sig.parameters[i] for i in sig.parameters if i != 'xp']) + + if wrapped_f.__doc__ is None: + wrapped_f.__doc__ = f"""\ +Array API compatibility wrapper for {f.__name__}. + +See the corresponding documentation in NumPy/CuPy and/or the array API +specification for more details. + +""" + wrapped_f.__signature__ = new_sig + return wrapped_f + + return inner diff --git a/array_api_compat/common/__init__.py b/array_api_compat/common/__init__.py new file mode 100644 index 00000000..ce3f44dd --- /dev/null +++ b/array_api_compat/common/__init__.py @@ -0,0 +1 @@ +from ._helpers import * diff --git a/numpy_array_api_compat/_aliases.py b/array_api_compat/common/_aliases.py similarity index 53% rename from numpy_array_api_compat/_aliases.py rename to array_api_compat/common/_aliases.py index 0223c068..ec8bd6e1 100644 --- a/numpy_array_api_compat/_aliases.py +++ b/array_api_compat/common/_aliases.py @@ -6,30 +6,164 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from typing import Optional, Tuple, Union + from typing import Optional, Tuple, Union, List from ._typing import ndarray, Device, Dtype, NestedSequence, SupportsBufferProtocol from typing import NamedTuple +from types import ModuleType -import numpy as np +from ._helpers import _check_device, _is_numpy_array, get_namespace # Basic renames -acos = np.arccos -acosh = np.arccosh -asin = np.arcsin -asinh = np.arcsinh -atan = np.arctan -atan2 = np.arctan2 -atanh = np.arctanh -bitwise_left_shift = np.left_shift -bitwise_invert = np.invert -bitwise_right_shift = np.right_shift -bool = np.bool_ -concat = np.concatenate -pow = np.power +def acos(x, /, xp): + return xp.arccos(x) + +def acosh(x, /, xp): + return xp.arccosh(x) + +def asin(x, /, xp): + return xp.arcsin(x) + +def asinh(x, /, xp): + return xp.arcsinh(x) + +def atan(x, /, xp): + return xp.arctan(x) + +def atan2(x1, x2, /, xp): + return xp.arctan2(x1, x2) + +def atanh(x, /, xp): + return xp.arctanh(x) + +def bitwise_left_shift(x1, x2, /, xp): + return xp.left_shift(x1, x2) + +def bitwise_invert(x, /, xp): + return xp.invert(x) + +def bitwise_right_shift(x1, x2, /, xp): + return xp.right_shift(x1, x2) + +def concat(arrays: Union[Tuple[ndarray, ...], List[ndarray]], /, xp, *, axis: Optional[int] = 0) -> ndarray: + return xp.concatenate(arrays, axis=axis) + +def pow(x1, x2, /, xp): + return xp.power(x1, x2) # These functions are modified from the NumPy versions. +def arange( + start: Union[int, float], + /, + stop: Optional[Union[int, float]] = None, + step: Union[int, float] = 1, + *, + xp, + dtype: Optional[Dtype] = None, + device: Optional[Device] = None, +) -> ndarray: + _check_device(xp, device) + return xp.arange(start, stop=stop, step=step, dtype=dtype) + +def empty( + shape: Union[int, Tuple[int, ...]], + xp, + *, + dtype: Optional[Dtype] = None, + device: Optional[Device] = None, +) -> ndarray: + _check_device(xp, device) + return xp.empty(shape, dtype=dtype) + +def empty_like( + x: ndarray, /, xp, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None +) -> ndarray: + _check_device(xp, device) + return xp.empty_like(x, dtype=dtype) + +def eye( + n_rows: int, + n_cols: Optional[int] = None, + /, + *, + xp, + k: int = 0, + dtype: Optional[Dtype] = None, + device: Optional[Device] = None, +) -> ndarray: + _check_device(xp, device) + return xp.eye(n_rows, M=n_cols, k=k, dtype=dtype) + +def full( + shape: Union[int, Tuple[int, ...]], + fill_value: Union[int, float], + xp, + *, + dtype: Optional[Dtype] = None, + device: Optional[Device] = None, +) -> ndarray: + _check_device(xp, device) + return xp.full(shape, fill_value, dtype=dtype) + +def full_like( + x: ndarray, + /, + fill_value: Union[int, float], + *, + xp, + dtype: Optional[Dtype] = None, + device: Optional[Device] = None, +) -> ndarray: + _check_device(xp, device) + return xp.full_like(x, fill_value, dtype=dtype) + +def linspace( + start: Union[int, float], + stop: Union[int, float], + /, + num: int, + *, + xp, + dtype: Optional[Dtype] = None, + device: Optional[Device] = None, + endpoint: bool = True, +) -> ndarray: + _check_device(xp, device) + return xp.linspace(start, stop, num, dtype=dtype, endpoint=endpoint) + +def ones( + shape: Union[int, Tuple[int, ...]], + xp, + *, + dtype: Optional[Dtype] = None, + device: Optional[Device] = None, +) -> ndarray: + _check_device(xp, device) + return xp.ones(shape, dtype=dtype) + +def ones_like( + x: ndarray, /, xp, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None +) -> ndarray: + _check_device(xp, device) + return xp.ones_like(x, dtype=dtype) + +def zeros( + shape: Union[int, Tuple[int, ...]], + xp, + *, + dtype: Optional[Dtype] = None, + device: Optional[Device] = None, +) -> ndarray: + _check_device(xp, device) + return xp.zeros(shape, dtype=dtype) + +def zeros_like( + x: ndarray, /, xp, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None +) -> ndarray: + _check_device(xp, device) + return xp.zeros_like(x, dtype=dtype) + # np.unique() is split into four functions in the array API: # unique_all, unique_counts, unique_inverse, and unique_values (this is done # to remove polymorphic return types). @@ -53,8 +187,8 @@ class UniqueInverseResult(NamedTuple): inverse_indices: ndarray -def unique_all(x: ndarray, /) -> UniqueAllResult: - values, indices, inverse_indices, counts = np.unique( +def unique_all(x: ndarray, /, xp) -> UniqueAllResult: + values, indices, inverse_indices, counts = xp.unique( x, return_counts=True, return_index=True, @@ -72,8 +206,8 @@ def unique_all(x: ndarray, /) -> UniqueAllResult: ) -def unique_counts(x: ndarray, /) -> UniqueCountsResult: - res = np.unique( +def unique_counts(x: ndarray, /, xp) -> UniqueCountsResult: + res = xp.unique( x, return_counts=True, return_index=False, @@ -84,22 +218,22 @@ def unique_counts(x: ndarray, /) -> UniqueCountsResult: return UniqueCountsResult(*res) -def unique_inverse(x: ndarray, /) -> UniqueInverseResult: - values, inverse_indices = np.unique( +def unique_inverse(x: ndarray, /, xp) -> UniqueInverseResult: + values, inverse_indices = xp.unique( x, return_counts=False, return_index=False, return_inverse=True, equal_nan=False, ) - # np.unique() flattens inverse indices, but they need to share x's shape + # xp.unique() flattens inverse indices, but they need to share x's shape # See https://github.com/numpy/numpy/issues/20638 inverse_indices = inverse_indices.reshape(x.shape) return UniqueInverseResult(values, inverse_indices) -def unique_values(x: ndarray, /) -> ndarray: - return np.unique( +def unique_values(x: ndarray, /, xp) -> ndarray: + return xp.unique( x, return_counts=False, return_index=False, @@ -117,35 +251,33 @@ def astype(x: ndarray, dtype: Dtype, /, *, copy: bool = True) -> ndarray: def std( x: ndarray, /, + xp, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, correction: Union[int, float] = 0.0, # correction instead of ddof keepdims: bool = False, ) -> ndarray: - return np.std(x, axis=axis, ddof=correction, keepdims=keepdims) + return xp.std(x, axis=axis, ddof=correction, keepdims=keepdims) def var( x: ndarray, /, + xp, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, correction: Union[int, float] = 0.0, # correction instead of ddof keepdims: bool = False, ) -> ndarray: - return np.var(x, axis=axis, ddof=correction, keepdims=keepdims) + return xp.var(x, axis=axis, ddof=correction, keepdims=keepdims) # Unlike transpose(), the axes argument to permute_dims() is required. -def permute_dims(x: ndarray, /, axes: Tuple[int, ...]) -> ndarray: - return np.transpose(x, axes) +def permute_dims(x: ndarray, /, axes: Tuple[int, ...], xp) -> ndarray: + return xp.transpose(x, axes) # Creation functions add the device keyword (which does nothing for NumPy) -def _check_device(device): - if device not in ["cpu", None]: - raise ValueError(f"Unsupported device {device!r}") - # asarray also adds the copy keyword -def asarray( +def _asarray( obj: Union[ ndarray, bool, @@ -158,148 +290,75 @@ def asarray( *, dtype: Optional[Dtype] = None, device: Optional[Device] = None, - copy: Optional[Union[bool, np._CopyMode]] = None, + copy: "Optional[Union[bool, np._CopyMode]]" = None, + namespace = None, ) -> ndarray: - _check_device(device) - if copy in (False, np._CopyMode.IF_NEEDED): - # copy=False is not yet implemented in np.asarray + """ + Array API compatibility wrapper for asarray(). + + See the corresponding documentation in NumPy/CuPy and/or the array API + specification for more details. + + """ + if namespace is None: + try: + xp = get_namespace(obj, _use_compat=False) + except ValueError: + # TODO: What about lists of arrays? + raise ValueError("A namespace must be specified for asarray() with non-array input") + elif isinstance(namespace, ModuleType): + xp = namespace + elif namespace == 'numpy': + import numpy as xp + elif namespace == 'cupy': + import cupy as xp + else: + raise ValueError("Unrecognized namespace argument to asarray()") + + _check_device(xp, device) + if _is_numpy_array(obj): + import numpy as np + COPY_FALSE = (False, np._CopyMode.IF_NEEDED) + COPY_TRUE = (True, np._CopyMode.ALWAYS) + else: + COPY_FALSE = (False,) + COPY_TRUE = (True,) + if copy in COPY_FALSE: + # copy=False is not yet implemented in xp.asarray raise NotImplementedError("copy=False is not yet implemented") - if isinstance(obj, np.ndarray): + if isinstance(obj, xp.ndarray): if dtype is not None and obj.dtype != dtype: copy = True - if copy in (True, np._CopyMode.ALWAYS): - return np.array(obj, copy=True, dtype=dtype) + if copy in COPY_TRUE: + return xp.array(obj, copy=True, dtype=dtype) return obj - return np.asarray(obj, dtype=dtype) - -def arange( - start: Union[int, float], - /, - stop: Optional[Union[int, float]] = None, - step: Union[int, float] = 1, - *, - dtype: Optional[Dtype] = None, - device: Optional[Device] = None, -) -> ndarray: - _check_device(device) - return np.arange(start, stop=stop, step=step, dtype=dtype) - -def empty( - shape: Union[int, Tuple[int, ...]], - *, - dtype: Optional[Dtype] = None, - device: Optional[Device] = None, -) -> ndarray: - _check_device(device) - return np.empty(shape, dtype=dtype) - -def empty_like( - x: ndarray, /, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None -) -> ndarray: - _check_device(device) - return np.empty_like(x, dtype=dtype) - -def eye( - n_rows: int, - n_cols: Optional[int] = None, - /, - *, - k: int = 0, - dtype: Optional[Dtype] = None, - device: Optional[Device] = None, -) -> ndarray: - _check_device(device) - return np.eye(n_rows, M=n_cols, k=k, dtype=dtype) - -def full( - shape: Union[int, Tuple[int, ...]], - fill_value: Union[int, float], - *, - dtype: Optional[Dtype] = None, - device: Optional[Device] = None, -) -> ndarray: - _check_device(device) - return np.full(shape, fill_value, dtype=dtype) - -def full_like( - x: ndarray, - /, - fill_value: Union[int, float], - *, - dtype: Optional[Dtype] = None, - device: Optional[Device] = None, -) -> ndarray: - _check_device(device) - return np.full_like(x, fill_value, dtype=dtype) - -def linspace( - start: Union[int, float], - stop: Union[int, float], - /, - num: int, - *, - dtype: Optional[Dtype] = None, - device: Optional[Device] = None, - endpoint: bool = True, -) -> ndarray: - _check_device(device) - return np.linspace(start, stop, num, dtype=dtype, endpoint=endpoint) + return xp.asarray(obj, dtype=dtype) -def ones( - shape: Union[int, Tuple[int, ...]], - *, - dtype: Optional[Dtype] = None, - device: Optional[Device] = None, -) -> ndarray: - _check_device(device) - return np.ones(shape, dtype=dtype) - -def ones_like( - x: ndarray, /, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None -) -> ndarray: - _check_device(device) - return np.ones_like(x, dtype=dtype) - -def zeros( - shape: Union[int, Tuple[int, ...]], - *, - dtype: Optional[Dtype] = None, - device: Optional[Device] = None, -) -> ndarray: - _check_device(device) - return np.zeros(shape, dtype=dtype) - -def zeros_like( - x: ndarray, /, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None -) -> ndarray: - _check_device(device) - return np.zeros_like(x, dtype=dtype) - -# np.reshape calls the keyword argument 'newshape' instead of 'shape' -def reshape(x: ndarray, /, shape: Tuple[int, ...], copy: Optional[bool] = None) -> ndarray: +# xp.reshape calls the keyword argument 'newshape' instead of 'shape' +def reshape(x: ndarray, /, shape: Tuple[int, ...], xp, copy: Optional[bool] = None) -> ndarray: if copy is True: x = x.copy() elif copy is False: x.shape = shape return x - return np.reshape(x, shape) + return xp.reshape(x, shape) # The descending keyword is new in sort and argsort, and 'kind' replaced with # 'stable' def argsort( - x: ndarray, /, *, axis: int = -1, descending: bool = False, stable: bool = True + x: ndarray, /, xp, *, axis: int = -1, descending: bool = False, stable: bool = True ) -> ndarray: # Note: this keyword argument is different, and the default is different. kind = "stable" if stable else "quicksort" if not descending: - res = np.argsort(x, axis=axis, kind=kind) + res = xp.argsort(x, axis=axis, kind=kind) else: # As NumPy has no native descending sort, we imitate it here. Note that - # simply flipping the results of np.argsort(x, ...) would not + # simply flipping the results of xp.argsort(x, ...) would not # respect the relative order like it would in native descending sorts. - res = np.flip( - np.argsort(np.flip(x, axis=axis), axis=axis, kind=kind), + res = xp.flip( + xp.argsort(xp.flip(x, axis=axis), axis=axis, kind=kind), axis=axis, ) # Rely on flip()/argsort() to validate axis @@ -309,67 +368,64 @@ def argsort( return res def sort( - x: ndarray, /, *, axis: int = -1, descending: bool = False, stable: bool = True + x: ndarray, /, xp, *, axis: int = -1, descending: bool = False, stable: bool = True ) -> ndarray: # Note: this keyword argument is different, and the default is different. kind = "stable" if stable else "quicksort" - res = np.sort(x, axis=axis, kind=kind) + res = xp.sort(x, axis=axis, kind=kind) if descending: - res = np.flip(res, axis=axis) + res = xp.flip(res, axis=axis) return res # sum() and prod() should always upcast when dtype=None def sum( x: ndarray, /, + xp, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Optional[Dtype] = None, keepdims: bool = False, ) -> ndarray: - # `np.sum` already upcasts integers, but not floats - if dtype is None and x.dtype == np.float32: - dtype = np.float64 - return np.sum(x, axis=axis, dtype=dtype, keepdims=keepdims) + # `xp.sum` already upcasts integers, but not floats + if dtype is None and x.dtype == xp.float32: + dtype = xp.float64 + return xp.sum(x, axis=axis, dtype=dtype, keepdims=keepdims) def prod( x: ndarray, /, + xp, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Optional[Dtype] = None, keepdims: bool = False, ) -> ndarray: - if dtype is None and x.dtype == np.float32: - dtype = np.float64 - return np.prod(x, dtype=dtype, axis=axis, keepdims=keepdims) + if dtype is None and x.dtype == xp.float32: + dtype = xp.float64 + return xp.prod(x, dtype=dtype, axis=axis, keepdims=keepdims) # ceil, floor, and trunc return integers for integer inputs -def ceil(x: ndarray, /) -> ndarray: - if np.issubdtype(x.dtype, np.integer): +def ceil(x: ndarray, /, xp) -> ndarray: + if xp.issubdtype(x.dtype, xp.integer): return x - return np.ceil(x) + return xp.ceil(x) -def floor(x: ndarray, /) -> ndarray: - if np.issubdtype(x.dtype, np.integer): +def floor(x: ndarray, /, xp) -> ndarray: + if xp.issubdtype(x.dtype, xp.integer): return x - return np.floor(x) + return xp.floor(x) -def trunc(x: ndarray, /) -> ndarray: - if np.issubdtype(x.dtype, np.integer): +def trunc(x: ndarray, /, xp) -> ndarray: + if xp.issubdtype(x.dtype, xp.integer): return x - return np.trunc(x) - -# from numpy import * doesn't overwrite these builtin names -from numpy import abs, max, min, round + return xp.trunc(x) __all__ = ['acos', 'acosh', 'asin', 'asinh', 'atan', 'atan2', 'atanh', 'bitwise_left_shift', 'bitwise_invert', 'bitwise_right_shift', - 'bool', 'concat', 'pow', 'UniqueAllResult', 'UniqueCountsResult', + 'concat', 'pow', 'UniqueAllResult', 'UniqueCountsResult', 'UniqueInverseResult', 'unique_all', 'unique_counts', - 'unique_inverse', 'unique_values', 'astype', 'abs', 'max', 'min', - 'round', 'std', 'var', 'permute_dims', 'asarray', 'arange', - 'empty', 'empty_like', 'eye', 'full', 'full_like', 'linspace', - 'ones', 'ones_like', 'zeros', 'zeros_like', 'reshape', 'argsort', - 'sort', 'sum', 'prod', 'ceil', 'floor', 'trunc'] + 'unique_inverse', 'unique_values', 'astype', 'std', 'var', + 'permute_dims', 'reshape', 'argsort', 'sort', 'sum', 'prod', + 'ceil', 'floor', 'trunc'] diff --git a/array_api_compat/common/_helpers.py b/array_api_compat/common/_helpers.py new file mode 100644 index 00000000..a1310b1c --- /dev/null +++ b/array_api_compat/common/_helpers.py @@ -0,0 +1,175 @@ +""" +Various helper functions which are not part of the spec. + +Functions which start with an underscore are for internal use only but helpers +that are in __all__ are intended as additional helper functions for use by end +users of the compat library. +""" +from __future__ import annotations + +import sys + +def _is_numpy_array(x): + # Avoid importing NumPy if it isn't already + if 'numpy' not in sys.modules: + return False + + import numpy as np + + # TODO: Should we reject ndarray subclasses? + return isinstance(x, (np.ndarray, np.generic)) + +def _is_cupy_array(x): + # Avoid importing NumPy if it isn't already + if 'cupy' not in sys.modules: + return False + + import cupy as cp + + # TODO: Should we reject ndarray subclasses? + return isinstance(x, (cp.ndarray, cp.generic)) + +def is_array_api_obj(x): + """ + Check if x is an array API compatible array object. + """ + return _is_numpy_array(x) or _is_cupy_array(x) or hasattr(x, '__array_namespace__') + +def get_namespace(*xs, _use_compat=True): + """ + Get the array API compatible namespace for the arrays `xs`. + + `xs` should contain one or more arrays. + """ + namespaces = set() + for x in xs: + if isinstance(x, (tuple, list)): + namespaces.add(get_namespace(*x, _use_compat=_use_compat)) + elif hasattr(x, '__array_namespace__'): + namespaces.add(x.__array_namespace__) + elif _is_numpy_array(x): + if _use_compat: + from .. import numpy as numpy_namespace + namespaces.add(numpy_namespace) + else: + import numpy as np + namespaces.add(np) + elif _is_cupy_array(x): + if _use_compat: + from .. import cupy as cupy_namespace + namespaces.add(cupy_namespace) + else: + import cupy as cp + namespaces.add(cp) + else: + # TODO: Support Python scalars? + raise ValueError("The input is not a supported array type") + + if not namespaces: + raise ValueError("Unrecognized array input") + + if len(namespaces) != 1: + raise ValueError(f"Multiple namespaces for array inputs: {namespaces}") + + xp, = namespaces + + return xp + + +def _check_device(xp, device): + if xp == sys.modules.get('numpy'): + if device not in ["cpu", None]: + raise ValueError(f"Unsupported device for NumPy: {device!r}") + +# device() is not on numpy.ndarray and and to_device() is not on numpy.ndarray +# or cupy.ndarray. They are not included in array objects of this library +# because this library just reuses the respective ndarray classes without +# wrapping or subclassing them. These helper functions can be used instead of +# the wrapper functions for libraries that need to support both NumPy/CuPy and +# other libraries that use devices. +def device(x: "Array", /) -> "Device": + """ + Hardware device the array data resides on. + + Parameters + ---------- + x: array + array instance from NumPy or an array API compatible library. + + Returns + ------- + out: device + a ``device`` object (see the "Device Support" section of the array API specification). + """ + if _is_numpy_array(x): + return "cpu" + return x.device + +# Based on cupy.array_api.Array.to_device +def _cupy_to_device(x, device, /, stream=None): + import cupy as cp + from cupy.cuda import Device as _Device + from cupy.cuda import stream as stream_module + from cupy_backends.cuda.api import runtime + + if device == x.device: + return x + elif not isinstance(device, _Device): + raise ValueError(f"Unsupported device {device!r}") + else: + # see cupy/cupy#5985 for the reason how we handle device/stream here + prev_device = runtime.getDevice() + prev_stream: stream_module.Stream = None + if stream is not None: + prev_stream = stream_module.get_current_stream() + # stream can be an int as specified in __dlpack__, or a CuPy stream + if isinstance(stream, int): + stream = cp.cuda.ExternalStream(stream) + elif isinstance(stream, cp.cuda.Stream): + pass + else: + raise ValueError('the input stream is not recognized') + stream.use() + try: + runtime.setDevice(device.id) + arr = x.copy() + finally: + runtime.setDevice(prev_device) + if stream is not None: + prev_stream.use() + return arr + +def to_device(x: "Array", device: "Device", /, *, stream: Optional[Union[int, Any]] = None) -> "Array": + """ + Copy the array from the device on which it currently resides to the specified ``device``. + + Parameters + ---------- + x: array + array instance from NumPy or an array API compatible library. + device: device + a ``device`` object (see the "Device Support" section of the array API specification). + stream: Optional[Union[int, Any]] + stream object to use during copy. In addition to the types supported in ``array.__dlpack__``, implementations may choose to support any library-specific stream object with the caveat that any code using such an object would not be portable. + + Returns + ------- + out: array + an array with the same data and data type as ``x`` and located on the specified ``device``. + + .. note:: + If ``stream`` is given, the copy operation should be enqueued on the provided ``stream``; otherwise, the copy operation should be enqueued on the default stream/queue. Whether the copy is performed synchronously or asynchronously is implementation-dependent. Accordingly, if synchronization is required to guarantee data safety, this must be clearly explained in a conforming library's documentation. + """ + if _is_numpy_array(x): + if stream is not None: + raise ValueError("The stream argument to to_device() is not supported") + if device == 'cpu': + return x + raise ValueError(f"Unsupported device {device!r}") + elif _is_cupy_array(x): + # cupy does not yet have to_device + return _cupy_to_device(x, device, stream=stream) + + return x.to_device(device, stream=stream) + +__all__ = ['is_array_api_obj', 'get_namespace', 'device', 'to_device'] diff --git a/array_api_compat/common/_linalg.py b/array_api_compat/common/_linalg.py new file mode 100644 index 00000000..f3bec324 --- /dev/null +++ b/array_api_compat/common/_linalg.py @@ -0,0 +1,165 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, NamedTuple +if TYPE_CHECKING: + from typing import Literal, Optional, Sequence, Tuple, Union + from ._typing import ndarray + +from numpy.core.numeric import normalize_axis_tuple + +from .._internal import get_xp + +# These are in the main NumPy namespace but not in numpy.linalg +def cross(x1: ndarray, x2: ndarray, /, xp, *, axis: int = -1) -> ndarray: + return xp.cross(x1, x2, axis=axis) + +def matmul(x1: ndarray, x2: ndarray, /, xp) -> ndarray: + return xp.matmul(x1, x2) + +def outer(x1: ndarray, x2: ndarray, /, xp) -> ndarray: + return xp.outer(x1, x2) + +def tensordot(x1: ndarray, x2: ndarray, /, xp, *, axes: Union[int, Tuple[Sequence[int], Sequence[int]]] = 2) -> ndarray: + return xp.tensordot(x1, x2, axes=axes) + +class EighResult(NamedTuple): + eigenvalues: ndarray + eigenvectors: ndarray + +class QRResult(NamedTuple): + Q: ndarray + R: ndarray + +class SlogdetResult(NamedTuple): + sign: ndarray + logabsdet: ndarray + +class SVDResult(NamedTuple): + U: ndarray + S: ndarray + Vh: ndarray + +# These functions are the same as their NumPy counterparts except they return +# a namedtuple. +def eigh(x: ndarray, /, xp) -> EighResult: + return EighResult(*xp.linalg.eigh(x)) + +def qr(x: ndarray, /, xp, *, mode: Literal['reduced', 'complete'] = 'reduced') -> QRResult: + return QRResult(*xp.linalg.qr(x, mode=mode)) + +def slogdet(x: ndarray, /, xp) -> SlogdetResult: + return SlogdetResult(*xp.linalg.slogdet(x)) + +def svd(x: ndarray, /, xp, *, full_matrices: bool = True) -> SVDResult: + return SVDResult(*xp.linalg.svd(x, full_matrices=full_matrices)) + +# These functions have additional keyword arguments + +# The upper keyword argument is new from NumPy +def cholesky(x: ndarray, /, xp, *, upper: bool = False) -> ndarray: + L = xp.linalg.cholesky(x) + if upper: + return get_xp(xp)(matrix_transpose)(L) + return L + +# The rtol keyword argument of matrix_rank() and pinv() is new from NumPy. +# Note that it has a different semantic meaning from tol and rcond. +def matrix_rank(x: ndarray, /, xp, *, rtol: Optional[Union[float, ndarray]] = None) -> ndarray: + # this is different from xp.linalg.matrix_rank, which supports 1 + # dimensional arrays. + if x.ndim < 2: + raise xp.linalg.LinAlgError("1-dimensional array given. Array must be at least two-dimensional") + S = xp.linalg.svd(x, compute_uv=False) + if rtol is None: + tol = S.max(axis=-1, keepdims=True) * max(x.shape[-2:]) * xp.finfo(S.dtype).eps + else: + # this is different from xp.linalg.matrix_rank, which does not + # multiply the tolerance by the largest singular value. + tol = S.max(axis=-1, keepdims=True)*xp.asarray(rtol)[..., xp.newaxis] + return xp.count_nonzero(S > tol, axis=-1) + +def pinv(x: ndarray, /, xp, *, rtol: Optional[Union[float, ndarray]] = None) -> ndarray: + # this is different from xp.linalg.pinv, which does not multiply the + # default tolerance by max(M, N). + if rtol is None: + rtol = max(x.shape[-2:]) * xp.finfo(x.dtype).eps + return xp.linalg.pinv(x, rcond=rtol) + +# These functions are new in the array API spec + +def matrix_norm(x: ndarray, /, xp, *, keepdims: bool = False, ord: Optional[Union[int, float, Literal['fro', 'nuc']]] = 'fro') -> ndarray: + return xp.linalg.norm(x, axis=(-2, -1), keepdims=keepdims, ord=ord) + +# Unlike transpose, matrix_transpose only transposes the last two axes. +def matrix_transpose(x: ndarray, /, xp) -> ndarray: + if x.ndim < 2: + raise ValueError("x must be at least 2-dimensional for matrix_transpose") + return xp.swapaxes(x, -1, -2) + +# svdvals is not in NumPy (but it is in SciPy). It is equivalent to +# xp.linalg.svd(compute_uv=False). +def svdvals(x: ndarray, /, xp) -> Union[ndarray, Tuple[ndarray, ...]]: + return xp.linalg.svd(x, compute_uv=False) + +def vecdot(x1: ndarray, x2: ndarray, /, xp, *, axis: int = -1) -> ndarray: + ndim = max(x1.ndim, x2.ndim) + x1_shape = (1,)*(ndim - x1.ndim) + tuple(x1.shape) + x2_shape = (1,)*(ndim - x2.ndim) + tuple(x2.shape) + if x1_shape[axis] != x2_shape[axis]: + raise ValueError("x1 and x2 must have the same size along the given axis") + + x1_, x2_ = xp.broadcast_arrays(x1, x2) + x1_ = xp.moveaxis(x1_, axis, -1) + x2_ = xp.moveaxis(x2_, axis, -1) + + res = x1_[..., None, :] @ x2_[..., None] + return res[..., 0, 0] + +def vector_norm(x: ndarray, /, xp, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False, ord: Optional[Union[int, float]] = 2) -> ndarray: + # xp.linalg.norm tries to do a matrix norm whenever axis is a 2-tuple or + # when axis=None and the input is 2-D, so to force a vector norm, we make + # it so the input is 1-D (for axis=None), or reshape so that norm is done + # on a single dimension. + if axis is None: + # Note: xp.linalg.norm() doesn't handle 0-D arrays + x = x.ravel() + _axis = 0 + elif isinstance(axis, tuple): + # Note: The axis argument supports any number of axes, whereas + # xp.linalg.norm() only supports a single axis for vector norm. + normalized_axis = normalize_axis_tuple(axis, x.ndim) + rest = tuple(i for i in range(x.ndim) if i not in normalized_axis) + newshape = axis + rest + x = xp.transpose(x, newshape).reshape( + (xp.prod([x.shape[i] for i in axis], dtype=int), *[x.shape[i] for i in rest])) + _axis = 0 + else: + _axis = axis + + res = xp.linalg.norm(x, axis=_axis, ord=ord) + + if keepdims: + # We can't reuse xp.linalg.norm(keepdims) because of the reshape hacks + # above to avoid matrix norm logic. + shape = list(x.shape) + _axis = normalize_axis_tuple(range(x.ndim) if axis is None else axis, x.ndim) + for i in _axis: + shape[i] = 1 + res = xp.reshape(res, tuple(shape)) + + return res + +# xp.diagonal and xp.trace operate on the first two axes whereas these +# operates on the last two + +def diagonal(x: ndarray, /, xp, *, offset: int = 0) -> ndarray: + return xp.diagonal(x, offset=offset, axis1=-2, axis2=-1) + +def trace(x: ndarray, /, xp, *, offset: int = 0) -> ndarray: + return xp.asarray(xp.trace(x, offset=offset, axis1=-2, axis2=-1)) + +__all__ = ['cross', 'matmul', 'outer', 'tensordot', 'EighResult', + 'QRResult', 'SlogdetResult', 'SVDResult', 'eigh', 'qr', 'slogdet', + 'svd', 'cholesky', 'matrix_rank', 'pinv', 'matrix_norm', + 'matrix_transpose', 'svdvals', 'vecdot', 'vector_norm', 'diagonal', + 'trace'] diff --git a/array_api_compat/common/_typing.py b/array_api_compat/common/_typing.py new file mode 100644 index 00000000..3f178060 --- /dev/null +++ b/array_api_compat/common/_typing.py @@ -0,0 +1,20 @@ +from __future__ import annotations + +__all__ = [ + "NestedSequence", + "SupportsBufferProtocol", +] + +from typing import ( + Any, + TypeVar, + Protocol, +) + +_T_co = TypeVar("_T_co", covariant=True) + +class NestedSequence(Protocol[_T_co]): + def __getitem__(self, key: int, /) -> _T_co | NestedSequence[_T_co]: ... + def __len__(self, /) -> int: ... + +SupportsBufferProtocol = Any diff --git a/array_api_compat/cupy/__init__.py b/array_api_compat/cupy/__init__.py new file mode 100644 index 00000000..e4f43c13 --- /dev/null +++ b/array_api_compat/cupy/__init__.py @@ -0,0 +1,22 @@ +from cupy import * + +# from cupy import * doesn't overwrite these builtin names +from cupy import abs, max, min, round + +# These imports may overwrite names from the import * above. +from ._aliases import * + +# Don't know why, but we have to do an absolute import to import linalg. If we +# instead do +# +# from . import linalg +# +# It doesn't overwrite cupy.linalg from above. The import is generated +# dynamically so that the library can be vendored. +__import__(__package__ + '.linalg') + +from .linalg import matrix_transpose, vecdot + +from ..common._helpers import * + +__array_api_version__ = '2021.12' diff --git a/array_api_compat/cupy/_aliases.py b/array_api_compat/cupy/_aliases.py new file mode 100644 index 00000000..cbc89381 --- /dev/null +++ b/array_api_compat/cupy/_aliases.py @@ -0,0 +1,61 @@ +from __future__ import annotations + +from functools import partial + +from ..common import _aliases + +from .._internal import get_xp + +asarray = asarray_cupy = partial(_aliases._asarray, namespace='cupy') +asarray.__doc__ = _aliases._asarray.__doc__ +del partial + +import cupy as cp +bool = cp.bool_ + +acos = get_xp(cp)(_aliases.acos) +acosh = get_xp(cp)(_aliases.acosh) +asin = get_xp(cp)(_aliases.asin) +asinh = get_xp(cp)(_aliases.asinh) +atan = get_xp(cp)(_aliases.atan) +atan2 = get_xp(cp)(_aliases.atan2) +atanh = get_xp(cp)(_aliases.atanh) +bitwise_left_shift = get_xp(cp)(_aliases.bitwise_left_shift) +bitwise_invert = get_xp(cp)(_aliases.bitwise_invert) +bitwise_right_shift = get_xp(cp)(_aliases.bitwise_right_shift) +concat = get_xp(cp)(_aliases.concat) +pow = get_xp(cp)(_aliases.pow) +arange = get_xp(cp)(_aliases.arange) +empty = get_xp(cp)(_aliases.empty) +empty_like = get_xp(cp)(_aliases.empty_like) +eye = get_xp(cp)(_aliases.eye) +full = get_xp(cp)(_aliases.full) +full_like = get_xp(cp)(_aliases.full_like) +linspace = get_xp(cp)(_aliases.linspace) +ones = get_xp(cp)(_aliases.ones) +ones_like = get_xp(cp)(_aliases.ones_like) +zeros = get_xp(cp)(_aliases.zeros) +zeros_like = get_xp(cp)(_aliases.zeros_like) +UniqueAllResult = get_xp(cp)(_aliases.UniqueAllResult) +UniqueCountsResult = get_xp(cp)(_aliases.UniqueCountsResult) +UniqueInverseResult = get_xp(cp)(_aliases.UniqueInverseResult) +unique_all = get_xp(cp)(_aliases.unique_all) +unique_counts = get_xp(cp)(_aliases.unique_counts) +unique_inverse = get_xp(cp)(_aliases.unique_inverse) +unique_values = get_xp(cp)(_aliases.unique_values) +astype = _aliases.astype +std = get_xp(cp)(_aliases.std) +var = get_xp(cp)(_aliases.var) +permute_dims = get_xp(cp)(_aliases.permute_dims) +reshape = get_xp(cp)(_aliases.reshape) +argsort = get_xp(cp)(_aliases.argsort) +sort = get_xp(cp)(_aliases.sort) +sum = get_xp(cp)(_aliases.sum) +prod = get_xp(cp)(_aliases.prod) +ceil = get_xp(cp)(_aliases.ceil) +floor = get_xp(cp)(_aliases.floor) +trunc = get_xp(cp)(_aliases.trunc) + +__all__ = _aliases.__all__ + ['asarray', 'asarray_cupy', 'bool', 'arange', + 'empty', 'empty_like', 'eye', 'full', 'full_like', + 'linspace', 'ones', 'ones_like', 'zeros', 'zeros_like'] diff --git a/array_api_compat/cupy/_typing.py b/array_api_compat/cupy/_typing.py new file mode 100644 index 00000000..f3d9aab6 --- /dev/null +++ b/array_api_compat/cupy/_typing.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +__all__ = [ + "ndarray", + "Device", + "Dtype", +] + +import sys +from typing import ( + Union, + TYPE_CHECKING, +) + +from cupy import ( + ndarray, + dtype, + int8, + int16, + int32, + int64, + uint8, + uint16, + uint32, + uint64, + float32, + float64, +) + +from cupy.cuda.device import Device + +if TYPE_CHECKING or sys.version_info >= (3, 9): + Dtype = dtype[Union[ + int8, + int16, + int32, + int64, + uint8, + uint16, + uint32, + uint64, + float32, + float64, + ]] +else: + Dtype = dtype diff --git a/array_api_compat/cupy/linalg.py b/array_api_compat/cupy/linalg.py new file mode 100644 index 00000000..04c71dec --- /dev/null +++ b/array_api_compat/cupy/linalg.py @@ -0,0 +1,44 @@ +from cupy.linalg import * +# cupy.linalg doesn't have __all__. If it is added, replace this with +# +# from cupy.linalg import __all__ as linalg_all +_n = {} +exec('from cupy.linalg import *', _n) +del _n['__builtins__'] +linalg_all = list(_n) +del _n + +from ..common import _linalg +from .._internal import get_xp + +import cupy as cp + +cross = get_xp(cp)(_linalg.cross) +matmul = get_xp(cp)(_linalg.matmul) +outer = get_xp(cp)(_linalg.outer) +tensordot = get_xp(cp)(_linalg.tensordot) +EighResult = _linalg.EighResult +QRResult = _linalg.QRResult +SlogdetResult = _linalg.SlogdetResult +SVDResult = _linalg.SVDResult +eigh = get_xp(cp)(_linalg.eigh) +qr = get_xp(cp)(_linalg.qr) +slogdet = get_xp(cp)(_linalg.slogdet) +svd = get_xp(cp)(_linalg.svd) +cholesky = get_xp(cp)(_linalg.cholesky) +matrix_rank = get_xp(cp)(_linalg.matrix_rank) +pinv = get_xp(cp)(_linalg.pinv) +matrix_norm = get_xp(cp)(_linalg.matrix_norm) +matrix_transpose = get_xp(cp)(_linalg.matrix_transpose) +svdvals = get_xp(cp)(_linalg.svdvals) +vecdot = get_xp(cp)(_linalg.vecdot) +vector_norm = get_xp(cp)(_linalg.vector_norm) +diagonal = get_xp(cp)(_linalg.diagonal) +trace = get_xp(cp)(_linalg.trace) + +__all__ = linalg_all + _linalg.__all__ + +del get_xp +del cp +del linalg_all +del _linalg diff --git a/array_api_compat/numpy/__init__.py b/array_api_compat/numpy/__init__.py new file mode 100644 index 00000000..745367bc --- /dev/null +++ b/array_api_compat/numpy/__init__.py @@ -0,0 +1,22 @@ +from numpy import * + +# from numpy import * doesn't overwrite these builtin names +from numpy import abs, max, min, round + +# These imports may overwrite names from the import * above. +from ._aliases import * + +# Don't know why, but we have to do an absolute import to import linalg. If we +# instead do +# +# from . import linalg +# +# It doesn't overwrite np.linalg from above. The import is generated +# dynamically so that the library can be vendored. +__import__(__package__ + '.linalg') + +from .linalg import matrix_transpose, vecdot + +from ..common._helpers import * + +__array_api_version__ = '2021.12' diff --git a/array_api_compat/numpy/_aliases.py b/array_api_compat/numpy/_aliases.py new file mode 100644 index 00000000..e6ff2ee2 --- /dev/null +++ b/array_api_compat/numpy/_aliases.py @@ -0,0 +1,61 @@ +from __future__ import annotations + +from functools import partial + +from ..common import _aliases + +from .._internal import get_xp + +asarray = asarray_numpy = partial(_aliases._asarray, namespace='numpy') +asarray.__doc__ = _aliases._asarray.__doc__ +del partial + +import numpy as np +bool = np.bool_ + +acos = get_xp(np)(_aliases.acos) +acosh = get_xp(np)(_aliases.acosh) +asin = get_xp(np)(_aliases.asin) +asinh = get_xp(np)(_aliases.asinh) +atan = get_xp(np)(_aliases.atan) +atan2 = get_xp(np)(_aliases.atan2) +atanh = get_xp(np)(_aliases.atanh) +bitwise_left_shift = get_xp(np)(_aliases.bitwise_left_shift) +bitwise_invert = get_xp(np)(_aliases.bitwise_invert) +bitwise_right_shift = get_xp(np)(_aliases.bitwise_right_shift) +concat = get_xp(np)(_aliases.concat) +pow = get_xp(np)(_aliases.pow) +arange = get_xp(np)(_aliases.arange) +empty = get_xp(np)(_aliases.empty) +empty_like = get_xp(np)(_aliases.empty_like) +eye = get_xp(np)(_aliases.eye) +full = get_xp(np)(_aliases.full) +full_like = get_xp(np)(_aliases.full_like) +linspace = get_xp(np)(_aliases.linspace) +ones = get_xp(np)(_aliases.ones) +ones_like = get_xp(np)(_aliases.ones_like) +zeros = get_xp(np)(_aliases.zeros) +zeros_like = get_xp(np)(_aliases.zeros_like) +UniqueAllResult = get_xp(np)(_aliases.UniqueAllResult) +UniqueCountsResult = get_xp(np)(_aliases.UniqueCountsResult) +UniqueInverseResult = get_xp(np)(_aliases.UniqueInverseResult) +unique_all = get_xp(np)(_aliases.unique_all) +unique_counts = get_xp(np)(_aliases.unique_counts) +unique_inverse = get_xp(np)(_aliases.unique_inverse) +unique_values = get_xp(np)(_aliases.unique_values) +astype = _aliases.astype +std = get_xp(np)(_aliases.std) +var = get_xp(np)(_aliases.var) +permute_dims = get_xp(np)(_aliases.permute_dims) +reshape = get_xp(np)(_aliases.reshape) +argsort = get_xp(np)(_aliases.argsort) +sort = get_xp(np)(_aliases.sort) +sum = get_xp(np)(_aliases.sum) +prod = get_xp(np)(_aliases.prod) +ceil = get_xp(np)(_aliases.ceil) +floor = get_xp(np)(_aliases.floor) +trunc = get_xp(np)(_aliases.trunc) + +__all__ = _aliases.__all__ + ['asarray', 'asarray_numpy', 'bool', 'arange', + 'empty', 'empty_like', 'eye', 'full', 'full_like', + 'linspace', 'ones', 'ones_like', 'zeros', 'zeros_like'] diff --git a/numpy_array_api_compat/_typing.py b/array_api_compat/numpy/_typing.py similarity index 66% rename from numpy_array_api_compat/_typing.py rename to array_api_compat/numpy/_typing.py index f49868e1..c5ebb5ab 100644 --- a/numpy_array_api_compat/_typing.py +++ b/array_api_compat/numpy/_typing.py @@ -4,18 +4,13 @@ "ndarray", "Device", "Dtype", - "NestedSequence", - "SupportsBufferProtocol", ] import sys from typing import ( - Any, Literal, Union, TYPE_CHECKING, - TypeVar, - Protocol, ) from numpy import ( @@ -33,12 +28,6 @@ float64, ) -_T_co = TypeVar("_T_co", covariant=True) - -class NestedSequence(Protocol[_T_co]): - def __getitem__(self, key: int, /) -> _T_co | NestedSequence[_T_co]: ... - def __len__(self, /) -> int: ... - Device = Literal["cpu"] if TYPE_CHECKING or sys.version_info >= (3, 9): Dtype = dtype[Union[ @@ -55,5 +44,3 @@ def __len__(self, /) -> int: ... ]] else: Dtype = dtype - -SupportsBufferProtocol = Any diff --git a/array_api_compat/numpy/linalg.py b/array_api_compat/numpy/linalg.py new file mode 100644 index 00000000..ac04b055 --- /dev/null +++ b/array_api_compat/numpy/linalg.py @@ -0,0 +1,37 @@ +from numpy.linalg import * +from numpy.linalg import __all__ as linalg_all + +from ..common import _linalg +from .._internal import get_xp + +import numpy as np + +cross = get_xp(np)(_linalg.cross) +matmul = get_xp(np)(_linalg.matmul) +outer = get_xp(np)(_linalg.outer) +tensordot = get_xp(np)(_linalg.tensordot) +EighResult = _linalg.EighResult +QRResult = _linalg.QRResult +SlogdetResult = _linalg.SlogdetResult +SVDResult = _linalg.SVDResult +eigh = get_xp(np)(_linalg.eigh) +qr = get_xp(np)(_linalg.qr) +slogdet = get_xp(np)(_linalg.slogdet) +svd = get_xp(np)(_linalg.svd) +cholesky = get_xp(np)(_linalg.cholesky) +matrix_rank = get_xp(np)(_linalg.matrix_rank) +pinv = get_xp(np)(_linalg.pinv) +matrix_norm = get_xp(np)(_linalg.matrix_norm) +matrix_transpose = get_xp(np)(_linalg.matrix_transpose) +svdvals = get_xp(np)(_linalg.svdvals) +vecdot = get_xp(np)(_linalg.vecdot) +vector_norm = get_xp(np)(_linalg.vector_norm) +diagonal = get_xp(np)(_linalg.diagonal) +trace = get_xp(np)(_linalg.trace) + +__all__ = linalg_all + _linalg.__all__ + +del get_xp +del np +del linalg_all +del _linalg diff --git a/numpy_array_api_compat/__init__.py b/numpy_array_api_compat/__init__.py deleted file mode 100644 index 82f36e98..00000000 --- a/numpy_array_api_compat/__init__.py +++ /dev/null @@ -1,60 +0,0 @@ -""" -NumPy Array API compatibility library - -This is a small wrapper around NumPy that is compatible with the Array API -standard https://data-apis.org/array-api/latest/. See also NEP 47 -https://numpy.org/neps/nep-0047-array-api-standard.html. - -Unlike numpy.array_api, this is not a strict minimal implementation of the -Array API, but rather just an extension of the main NumPy namespace with -changes needed to be compliant with the Array API. See -https://numpy.org/doc/stable/reference/array_api.html for a full list of -changes. In particular, unlike numpy.array_api, this package does not use a -separate Array object, but rather just uses numpy.ndarray directly. - -Library authors using the Array API may wish to test against numpy.array_api -to ensure they are not using functionality outside of the standard, but prefer -this implementation for the default when working with NumPy arrays. - -In addition, several helper functions are provided in this library which are -not part of the array API specification but which are useful for libraries -writing against the array API specification who wish to support NumPy and -other array API compatible libraries. - -Known differences from the Array API spec: - -- The array methods __array_namespace__, device, to_device, and mT are not - defined. This reuses np.ndarray and we don't want to monkeypatch or wrap it. - The helper functions device() and to_device() are provided to work around - these missing methods. x.mT can be replaced with - xp.linalg.matrix_transpose(x). - -- NumPy value-based casting for scalars will be in effect unless explicitly - disabled with the environment variable NPY_PROMOTION_STATE=weak or - np._set_promotion_state('weak') (requires NumPy 1.24 or newer, see NEP 50 - and https://github.com/numpy/numpy/issues/22341) - -- NumPy functions which are not wrapped may not have the same type annotations - as the spec. - -- NumPy functions which are not wrapped may not use positional-only arguments. - -""" - -from numpy import * - -# These imports may overwrite names from the import * above. -from ._aliases import * - -# Don't know why, but we have to do an absolute import to import linalg. If we -# instead do -# -# from . import linalg -# -# It doesn't overwrite np.linalg from above. The import is generated -# dynamically so that the library can be vendored. -__import__(__package__ + '.linalg') - -from .linalg import matrix_transpose, vecdot - -from ._helpers import * diff --git a/numpy_array_api_compat/_helpers.py b/numpy_array_api_compat/_helpers.py deleted file mode 100644 index 7b789bb6..00000000 --- a/numpy_array_api_compat/_helpers.py +++ /dev/null @@ -1,100 +0,0 @@ -""" -Various helper functions which are not part of the spec. -""" - -from __future__ import annotations - -import importlib -compat_namespace = importlib.import_module(__package__) - -import numpy as np - -def _is_numpy_array(x): - # TODO: Should we reject ndarray subclasses? - return isinstance(x, (np.ndarray, np.generic)) - -def is_array_api_obj(x): - """ - Check if x is an array API compatible array object. - """ - return _is_numpy_array(x) or hasattr(x, '__array_namespace__') - -def get_namespace(*xs): - """ - Get the array API compatible namespace for the arrays `xs`. - - `xs` should contain one or more arrays. - """ - namespaces = set() - for x in xs: - if hasattr(x, '__array_namespace__'): - namespaces.add(x.__array_namespace__) - elif _is_numpy_array(x): - namespaces.add(compat_namespace) - else: - # TODO: Support Python scalars? - raise ValueError("The input is not a supported array type") - - if not namespaces: - raise ValueError("Unrecognized array input") - - if len(namespaces) != 1: - raise ValueError(f"Multiple namespaces for array inputs: {namespaces}") - - xp, = namespaces - - return xp - -# device and to_device are not included in array object of this library -# because this library just reuses ndarray without wrapping or subclassing it. -# These helper functions can be used instead of the wrapper functions for -# libraries that need to support both NumPy and other libraries that use devices. -def device(x: "Array", /) -> "Device": - """ - Hardware device the array data resides on. - - Parameters - ---------- - x: array - array instance from NumPy or an array API compatible library. - - Returns - ------- - out: device - a ``device`` object (see the "Device Support" section of the array API specification). - """ - if _is_numpy_array(x): - return "cpu" - return x.device - -def to_device(x: "Array", device: "Device", /, *, stream: Optional[Union[int, Any]] = None) -> "Array": - """ - Copy the array from the device on which it currently resides to the specified ``device``. - - Parameters - ---------- - x: array - array instance from NumPy or an array API compatible library. - device: device - a ``device`` object (see the "Device Support" section of the array API specification). - stream: Optional[Union[int, Any]] - stream object to use during copy. In addition to the types supported in ``array.__dlpack__``, implementations may choose to support any library-specific stream object with the caveat that any code using such an object would not be portable. - - Returns - ------- - out: array - an array with the same data and data type as ``x`` and located on the specified ``device``. - - .. note:: - If ``stream`` is given, the copy operation should be enqueued on the provided ``stream``; otherwise, the copy operation should be enqueued on the default stream/queue. Whether the copy is performed synchronously or asynchronously is implementation-dependent. Accordingly, if synchronization is required to guarantee data safety, this must be clearly explained in a conforming library's documentation. - """ - if _is_numpy_array(x): - if stream is not None: - raise ValueError("The stream argument to to_device() is not supported") - if device == 'cpu': - return x - raise ValueError(f"Unsupported device {device!r}") - - return x.to_device(device, stream=stream) - -__all__ = ['is_array_api_obj', 'get_namespace', 'device', 'to_device'] diff --git a/numpy_array_api_compat/linalg.py b/numpy_array_api_compat/linalg.py deleted file mode 100644 index ec3bc11d..00000000 --- a/numpy_array_api_compat/linalg.py +++ /dev/null @@ -1,157 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING, NamedTuple -if TYPE_CHECKING: - from typing import Literal, Optional, Tuple, Union - from numpy import ndarray - -import numpy as np -from numpy.core.numeric import normalize_axis_tuple - -from numpy.linalg import * -from numpy.linalg import __all__ as linalg_all - -# These are in the main NumPy namespace but not in numpy.linalg -from numpy import cross, matmul, outer, tensordot - -class EighResult(NamedTuple): - eigenvalues: ndarray - eigenvectors: ndarray - -class QRResult(NamedTuple): - Q: ndarray - R: ndarray - -class SlogdetResult(NamedTuple): - sign: ndarray - logabsdet: ndarray - -class SVDResult(NamedTuple): - U: ndarray - S: ndarray - Vh: ndarray - -# These functions are the same as their NumPy counterparts except they return -# a namedtuple. -def eigh(x: ndarray, /) -> EighResult: - return EighResult(*np.linalg.eigh(x)) - -def qr(x: ndarray, /, *, mode: Literal['reduced', 'complete'] = 'reduced') -> QRResult: - return QRResult(*np.linalg.qr(x, mode=mode)) - -def slogdet(x: ndarray, /) -> SlogdetResult: - return SlogdetResult(*np.linalg.slogdet(x)) - -def svd(x: ndarray, /, *, full_matrices: bool = True) -> SVDResult: - return SVDResult(*np.linalg.svd(x, full_matrices=full_matrices)) - -# These functions have additional keyword arguments - -# The upper keyword argument is new from NumPy -def cholesky(x: ndarray, /, *, upper: bool = False) -> ndarray: - L = np.linalg.cholesky(x) - if upper: - return matrix_transpose(L) - return L - -# The rtol keyword argument of matrix_rank() and pinv() is new from NumPy. -# Note that it has a different semantic meaning from tol and rcond. -def matrix_rank(x: ndarray, /, *, rtol: Optional[Union[float, ndarray]] = None) -> ndarray: - # this is different from np.linalg.matrix_rank, which supports 1 - # dimensional arrays. - if x.ndim < 2: - raise np.linalg.LinAlgError("1-dimensional array given. Array must be at least two-dimensional") - S = np.linalg.svd(x, compute_uv=False) - if rtol is None: - tol = S.max(axis=-1, keepdims=True) * max(x.shape[-2:]) * np.finfo(S.dtype).eps - else: - # this is different from np.linalg.matrix_rank, which does not - # multiply the tolerance by the largest singular value. - tol = S.max(axis=-1, keepdims=True)*np.asarray(rtol)[..., np.newaxis] - return np.count_nonzero(S > tol, axis=-1) - -def pinv(x: ndarray, /, *, rtol: Optional[Union[float, ndarray]] = None) -> ndarray: - # this is different from np.linalg.pinv, which does not multiply the - # default tolerance by max(M, N). - if rtol is None: - rtol = max(x.shape[-2:]) * np.finfo(x.dtype).eps - return np.linalg.pinv(x, rcond=rtol) - -# These functions are new in the array API spec - -def matrix_norm(x: ndarray, /, *, keepdims: bool = False, ord: Optional[Union[int, float, Literal['fro', 'nuc']]] = 'fro') -> ndarray: - return np.linalg.norm(x, axis=(-2, -1), keepdims=keepdims, ord=ord) - -# Unlike transpose, matrix_transpose only transposes the last two axes. -def matrix_transpose(x: ndarray, /) -> ndarray: - if x.ndim < 2: - raise ValueError("x must be at least 2-dimensional for matrix_transpose") - return np.swapaxes(x, -1, -2) - -# svdvals is not in NumPy (but it is in SciPy). It is equivalent to -# np.linalg.svd(compute_uv=False). -def svdvals(x: ndarray, /) -> Union[ndarray, Tuple[ndarray, ...]]: - return np.linalg.svd(x, compute_uv=False) - -def vecdot(x1: ndarray, x2: ndarray, /, *, axis: int = -1) -> ndarray: - ndim = max(x1.ndim, x2.ndim) - x1_shape = (1,)*(ndim - x1.ndim) + tuple(x1.shape) - x2_shape = (1,)*(ndim - x2.ndim) + tuple(x2.shape) - if x1_shape[axis] != x2_shape[axis]: - raise ValueError("x1 and x2 must have the same size along the given axis") - - x1_, x2_ = np.broadcast_arrays(x1, x2) - x1_ = np.moveaxis(x1_, axis, -1) - x2_ = np.moveaxis(x2_, axis, -1) - - res = x1_[..., None, :] @ x2_[..., None] - return res[..., 0, 0] - -def vector_norm(x: ndarray, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False, ord: Optional[Union[int, float]] = 2) -> ndarray: - # np.linalg.norm tries to do a matrix norm whenever axis is a 2-tuple or - # when axis=None and the input is 2-D, so to force a vector norm, we make - # it so the input is 1-D (for axis=None), or reshape so that norm is done - # on a single dimension. - if axis is None: - # Note: np.linalg.norm() doesn't handle 0-D arrays - x = x.ravel() - _axis = 0 - elif isinstance(axis, tuple): - # Note: The axis argument supports any number of axes, whereas - # np.linalg.norm() only supports a single axis for vector norm. - normalized_axis = normalize_axis_tuple(axis, x.ndim) - rest = tuple(i for i in range(x.ndim) if i not in normalized_axis) - newshape = axis + rest - x = np.transpose(x, newshape).reshape( - (np.prod([x.shape[i] for i in axis], dtype=int), *[x.shape[i] for i in rest])) - _axis = 0 - else: - _axis = axis - - res = np.linalg.norm(x, axis=_axis, ord=ord) - - if keepdims: - # We can't reuse np.linalg.norm(keepdims) because of the reshape hacks - # above to avoid matrix norm logic. - shape = list(x.shape) - _axis = normalize_axis_tuple(range(x.ndim) if axis is None else axis, x.ndim) - for i in _axis: - shape[i] = 1 - res = np.reshape(res, tuple(shape)) - - return res - -# np.diagonal and np.trace operate on the first two axes whereas these -# operates on the last two - -def diagonal(x: ndarray, /, *, offset: int = 0) -> ndarray: - return np.diagonal(x, offset=offset, axis1=-2, axis2=-1) - -def trace(x: ndarray, /, *, offset: int = 0) -> ndarray: - return np.asarray(np.trace(x, offset=offset, axis1=-2, axis2=-1)) - -__all__ = linalg_all.copy() -__all__ += ['cross', 'diagonal', 'matmul', 'cholesky', 'matrix_rank', 'pinv', - 'matrix_norm', 'matrix_transpose', 'outer', 'svdvals', - 'tensordot', 'trace', 'vecdot', 'vector_norm', 'EighResult', - 'QRResult', 'SlogdetResult', 'SVDResult'] diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 00000000..abc8f0f1 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1,7 @@ +""" +Basic tests for the compat library + +This only tests basic things like that vendoring works. The extensive tests +are done by the array API test suite https://github.com/data-apis/array-api-tests + +""" diff --git a/tests/test_vendoring.py b/tests/test_vendoring.py new file mode 100644 index 00000000..85f68626 --- /dev/null +++ b/tests/test_vendoring.py @@ -0,0 +1,15 @@ +from pytest import skip + +def test_vendoring_numpy(): + from vendor_test import uses_numpy + uses_numpy._test_numpy() + + +def test_vendoring_cupy(): + try: + import cupy + except ImportError: + skip("CuPy is not installed") + + from vendor_test import uses_cupy + uses_cupy._test_cupy() diff --git a/vendor_test/__init__.py b/vendor_test/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/vendor_test/uses_cupy.py b/vendor_test/uses_cupy.py new file mode 100644 index 00000000..97f710b9 --- /dev/null +++ b/vendor_test/uses_cupy.py @@ -0,0 +1,18 @@ +# Basic test that vendoring works + +from .vendored._compat import cupy as cp_compat + +import cupy as cp + +def _test_cupy(): + a = cp_compat.asarray([1., 2., 3.]) + b = cp_compat.arange(3, dtype=cp_compat.float32) + + # cp.pow does not exist. Update this to use something else if it is added + res = cp_compat.pow(a, b) + assert res.dtype == cp_compat.float64 == cp.float64 + assert isinstance(a, cp.ndarray) + assert isinstance(b, cp.ndarray) + assert isinstance(res, cp.ndarray) + + cp.testing.assert_allclose(res, [1., 2., 9.]) diff --git a/vendor_test/uses_numpy.py b/vendor_test/uses_numpy.py new file mode 100644 index 00000000..96f2c5ff --- /dev/null +++ b/vendor_test/uses_numpy.py @@ -0,0 +1,18 @@ +# Basic test that vendoring works + +from .vendored._compat import numpy as np_compat + +import numpy as np + +def _test_numpy(): + a = np_compat.asarray([1., 2., 3.]) + b = np_compat.arange(3, dtype=np_compat.float32) + + # np.pow does not exist. Update this to use something else if it is added + res = np_compat.pow(a, b) + assert res.dtype == np_compat.float64 == np.float64 + assert isinstance(a, np.ndarray) + assert isinstance(b, np.ndarray) + assert isinstance(res, np.ndarray) + + np.testing.assert_allclose(res, [1., 2., 9.]) diff --git a/vendor_test/vendored/__init__.py b/vendor_test/vendored/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/vendor_test/vendored/_compat b/vendor_test/vendored/_compat new file mode 120000 index 00000000..ba484f91 --- /dev/null +++ b/vendor_test/vendored/_compat @@ -0,0 +1 @@ +../../array_api_compat/ \ No newline at end of file