Package: array-api-compat
Array API compatibility libraryThis is a small wrapper around common array libraries that is compatible with
theArray API standard. Currently,
NumPy, CuPy, and PyTorch are supported. If you want support for other array
libraries, or if you encounter any issues, pleaseopen an
issue.Note that some of the functionality in this library is backwards incompatible
with the corresponding wrapped libraries. The end-goal is to eventually make
each array library itself fully compatible with the array API, but this
requires making backwards incompatible changes in many cases, so this will
take some time.Currently all libraries here are implemented against the2022.12
versionof the standard.Installarray-api-compatis available on bothPyPIpython -m pip install array-api-compatandConda-forgeconda install --channel conda-forge array-api-compatUsageThe typical usage of this library will be to get the corresponding array API
compliant namespace from the input arrays usingarray_namespace(), likedefyour_function(x,y):xp=array_api_compat.array_namespace(x,y)# Now use xp as the array library namespacereturnxp.mean(x,axis=0)+2*xp.std(y,axis=0)If you wish to have library-specific code-paths, you can import the
corresponding wrapped namespace for each library, likeimportarray_api_compat.numpyasnpimportarray_api_compat.cupyascpimportarray_api_compat.torchastorchEach will include all the functions from the normal NumPy/CuPy/PyTorch
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 the same array object from the wrapped library.Difference betweenarray_api_compatandnumpy.array_apinumpy.array_apiis a strict minimal implementation of the Array API (seeNEP 47). For
example,numpy.array_apidoes not include any functions that are not part of
the array API specification, and will explicitly disallow behaviors that are
not required by the spec (e.g.,cross-kind type
promotions).
(cupy.array_apiis similar tonumpy.array_api)array_api_compat, on the other hand, is just an extension of the
corresponding array library namespaces with changes needed to be compliant
with the array API. It includes all additional library functions not mentioned
in the spec, and allows any library behaviors not explicitly disallowed by it,
such as cross-kind casting.In particular, unlikenumpy.array_api, this package does not use a separateArrayobject, but rather just uses the corresponding array library array
objects (numpy.ndarray,cupy.ndarray,torch.Tensor, etc.) directly. This
is because those are the objects that are going to be passed as inputs to
functions by end users. This does mean that a few behaviors cannot be wrapped
(see below), but most of the array API functional, so this does not affect
most things.Array consuming library authors coding against the array API may wish to test
againstnumpy.array_apito ensure they are not using functionality outside
of the standard, but prefer this implementation for the default behavior for
end-users.Helper FunctionsIn addition to the wrapped library namespaces and functions in the array API
specification, there are several helper functions included here that aren't
part of the specification but which are useful for using the array API:is_array_api_obj(x): ReturnTrueifxis an array API compatible array
object.array_namespace(*xs): Get the corresponding array API namespace for the
arraysxs. For example, if the arrays are NumPy arrays, the returned
namespace will bearray_api_compat.numpy. Note that this function will
also work for namespaces that aren't supported by this compat library but
which do support the array API (i.e., arrays that have the__array_namespace__attribute).device(x): Equivalent tox.devicein the array API specification. Included becausenumpy.ndarraydoes not
include thedeviceattribute and this library does not wrap or extend the
array object. Note that for NumPy,device(x)is always"cpu".to_device(x, device, /, *, stream=None): Equivalent tox.to_device.
Included because neither NumPy's, CuPy's, nor PyTorch's array 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 CUDADeviceandStreamobjects. For PyTorch, this is the same asx.to(device)(thestreamargument is not supported in PyTorch).size(x): Equivalent tox.size,
i.e., the number of elements in the array. Included because PyTorch'sTensordefinessizeas a method which returns the shape, and this cannot
be wrapped because this compat library doesn't wrap or extend the array
objects.Known Differences from the Array API SpecificationThere are some known differences between this library and the array API
specification:NumPy and CuPyThe array methods__array_namespace__,device(for NumPy),to_device,
andmTare not defined. This reusesnp.ndarrayandcp.ndarrayand we
don't want to monkeypatch or wrap it. The helper functionsdevice()andto_device()are provided to work around these missing methods (see above).x.mTcan be replaced withxp.linalg.matrix_transpose(x).array_namespace(x)should be used instead ofx.__array_namespace__.Value-based casting for scalars will be in effect unless explicitly disabled
with the environment variableNPY_PROMOTION_STATE=weakornp._set_promotion_state('weak')(requires NumPy 1.24 or newer, seeNEP
50andhttps://github.com/numpy/numpy/issues/22341)asarray()does not supportcopy=False.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.The minimum supported NumPy version is 1.21. However, this older version of
NumPy has a few issues:unique_*will not compare nans as unequal.finfo()has nosmallest_normal.Nofrom_dlpackor__dlpack__.argmax()andargmin()do not havekeepdims.qr()doesn't support matrix stacks.asarray()doesn't supportcopy=True(as noted above,copy=Falseis not
supported even in the latest NumPy).Type promotion behavior will be value based for 0-D arrays (and there is noNPY_PROMOTION_STATE=weakto disable this).If any of these are an issue, it is recommended to bump your minimum NumPy
version.PyTorchLike NumPy/CuPy, we do not wrap thetorch.Tensorobject. It is missing the__array_namespace__andto_devicemethods, so the corresponding helper
functionsarray_namespace()andto_device()in this library should be
used instead (see above).Thex.sizeattribute ontorch.Tensoris a function that behaves
differently fromx.sizein the spec. Use thesize(x)helper function as a portable workaround (see
above).PyTorch does not have unsigned integer types other thanuint8, and no
attempt is made to implement them here.PyTorch has type promotion semantics that differ from the array API
specification for 0-D tensor objects. The array functions in this wrapper
library do work around this, but the operators on the Tensor object do not,
as no operators or methods on the Tensor object are modified. If this is a
concern, use the functional form instead of the operator form, e.g.,add(x, y)instead ofx + y.unique_all()is not implemented, due to the fact thattorch.uniquedoes not support
returning theindicesarray. The otherunique_*functions are implemented.Slices do not support negative steps.std()andvar()do not support floating-pointcorrection.Thestreamargument of theto_device()helper (see above) is not
supported.As with NumPy, type annotations and positional-only arguments may not
exactly match the spec for functions that are not wrapped at all.The minimum supported PyTorch version is 1.13.VendoringThis library supports vendoring as an installation method. To vendor the
library, simply copyarray_api_compatinto the appropriate place in the
library, likecp -R array_api_compat/ mylib/vendored/array_api_compatYou 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 NotesAs 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 returnnp.ndarrayfor NumPy andcp.ndarrayfor CuPy.Each namespace (array_api_compat.numpy,array_api_compat.cupy, andarray_api_compat.torch) is populated with the normal library namespace (likefrom numpy import *). Then specific functions are replaced with wrapped
variants.Since NumPy and CuPy are nearly identical in behavior, most wrapping logic can
be shared between them. Wrapped functions that have the same logic between
NumPy and CuPy are inarray_api_compat/common/.
These functions are defined like# In array_api_compat/common/_aliases.pydefacos(x,/,xp):returnxp.arccos(x)Thexpargument refers to the original array namespace (eithernumpyorcupy). Then in the specificarray_api_compat/numpy/andarray_api_compat/cupy/namespaces, the@get_xpdecorator is applied to
these functions, which automatically removes thexpargument from the
function signature and replaces it with the corresponding array library, like# In array_api_compat/numpy/_aliases.pyfrom..commonimport_aliasesimportnumpyasnpacos=get_xp(np)(_aliases.acos)Thisacosnow has the signatureacos(x, /)and callsnumpy.arccos.Similarly, for CuPy:# In array_api_compat/cupy/_aliases.pyfrom..commonimport_aliasesimportcupyascpacos=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.PyTorch uses a similar layout inarray_api_compat/torch/, but it differs
enough from NumPy/CuPy that very few common wrappers for those libraries are
reused.Seehttps://numpy.org/doc/stable/reference/array_api.htmlfor a full list of
changes from the base NumPy (the differences for CuPy are nearly identical). A
corresponding document does not yet exist for PyTorch, but you can examine the
various comments in theimplementationto see what functions and behaviors have been wrapped.