Array-API interfaces

[jorenham]

https://data-apis.org/array-api/latest/API_specification/index.html

goal

The goal is to provide Protocols that can be used to type e.g. scipy.fft.fftfreq (through its xp kwarg) and scipy.fft.fftshift (via the passed array).

the (big) problems

• The dtype spec only defines __eq__, which also exists in object. So it's impossible to (structurally) type a "dtype".
• In the same way, it's impossible to (structurally) type a "device".
• Shapes are defined as tuples of "something". So the best we can do is type the number of dimensions, and ignore the axis types.

This means that static type-checkers won't be able to see the difference between a boolean array on your CPU, and a complex-floating array on a quantum computer within the core of a neutron star in the Andromeda galaxy.

so now what?

We can still figure out the following types:

• Array + the dimensionality (because shape: tuple[?, ...] kinda tells us the ndim, but not the actual shape)
• Namespace, which is a protocollable module-type (often named xp), that could give us
  • xp.array: the specific Array type (but there's still no HKT after asking for it for over 6 years...), and
  • xp.(bool|((u?)int|float|complex)(8|16|32|64)) -i.e. the impossible-to-type DType types (🎩 🪄 🐰)
  • xp.__array_namespace_info__() -> Info (docs), an object with some other "magical" methods:
    • .default_device() -> DefaultDevice
    • .devices() -> list[Device]
    • .dtypes(*, device: Device?, kind: Kind?) - with Kind a literal string category, which allows us to "categorize" the e.g. floats, ints, and complexes (assuming the lib properly typed it), i.e. as type aliases.

problem solved: just type the untypeable

Note that we can also obtain the Info type via Array, since Array.__array_namespace__() -> Namespace[Info] s.t. Namespace[InfoT].__array_namespace_info__() -> InfoT (ignoring the versions).
So we can get all dtype types from either Array or Namespace, and we can either one from the other, and vice-versa (i.e. bidirectionally in both ways, and also the other way around).

[jorenham]

There are some "stubs" within the array-api repo already: https://github.com/data-apis/array-api/tree/2023.12/src/array_api_stubs

But they have one tiny issue; they don't make any sense at all, and are completely useless in their current state.
Sadly, it appears to me as if it'd be faster to rewrite them from scratch, then to attempt to fix it.

For example; the use of typing.TypeVar is almost always incorrect. It looks as if the authors believe that a TypeVar can be used as some sort of global placeholder for to-be-defined types 🤷🏻.

---

update

This comment seems to confirm my hunch:

"""
Types for type annotations used in the array API standard.

The type variables should be replaced with the actual types for a given library, e.g., for NumPy TypeVar('array') would be replaced with ndarray.
"""

from https://github.com/data-apis/array-api/blob/fbc9ddec2f0366df069901a1d99580b5791f8772/src/array_api_stubs/_2023_12/_types.py#L1C1-L6C4

[jorenham]

Typing Survey of Array-API {2022,2023}.12 libraries

The Array API specification is for the most part unfinished: There are no specifications at all for "dtypes" amd "devices". And no, device: Any | None isn't good enough: Any is the typing equivalent of # TODO.
So without a specification, there's no way to properly type it properly.

This is why an elaborate workaround is needed. But developing this workaround requires comparing all of the dtypes and device-types that the array-api-compatible libraries are using at the moment.
So I compared all of the (functional) array-api libraries that I could find, and figured out what the exact dtype- and device-types are that they use.

There are no common denominators.

The array-api initiative claimed this as their responsibility. And although I admire their intentions and appload what they have achieved; it currently is not possible to statically type anything array-api-related with a sufficient level of detail.

Either way, this survey gave me something to work with, which should be enough to develop some type aliases that could serve as a workaround for these libraries in question.

There's a good chance that this won't work for other array-api libraries.

Libraries

The following libraries have native support for the 2023.12 version:

• array_api_strict
• numpy (>=2.1.0)
• jax.numpy

Through the array-api-compat the following libraries also support the 2023.12 version:

• cupy
• dask.array
• torch

These don't support the 2023.12 version yet, and therefore have no __array_namespace_info__, so require manual
analysis:

• sparse (pydata/sparse#696)

Device types

The types of the __array_namespace_info__.devices():

Package	Device Type
array_api_strict	array_api_strict._array_object.Device
cupy	cupy.cuda.device.Device
dask.array	array_api_compat.common._helpers._dask_device | Literal['cpu']
jax.numpy	jax.lib.xla_client.Device
numpy	Literal['cpu']
torch	torch.device
sparse(< 0.16)	Literal['cpu']

TODO: Write as type alias (union of protocols)

DTypes

Package	DType
array_api_strict	array_api_strict._dtypes._DType
cupy	numpy.dtype
dask.array	numpy.dtype
jax.numpy	numpy.dtype
numpy	numpy.dtype
torch	torch.dtype
sparse (< 0.16)	numpy.dtype
sparse (≥ 0.16)	numpy.dtype |sparse.mlir_backend._dtypes.DType

Assuming that (only) numpy is available, then these the above DTypes can be aggregated as the following DType type alias.

from typing import Literal, Protocol
import numpy as np

type _DTypeNP = np.dtype[np.bool | np.number[_8Bit | _16Bit | _32Bit | _64Bit]]

# https://github.com/pytorch/pytorch/blob/v2.5.0/torch/_C/__init__.pyi.in#L181-L189
class _DTypeTorch(Protocol):
    @property
    def is_floating_point(self, /) -> bool: ...
    @property
    def is_complex(self, /) -> bool: ...
    @property
    def is_signed(self, /) -> bool: ...
    @property
    def itemsize(self, /) -> Literal[1, 2, 4, 8]: ...
    def to_real(self) -> dtype: ...
    def to_complex(self) -> dtype: ...

# https://github.com/data-apis/array-api-strict/blob/2.1/array_api_strict/_dtypes.py
class _DTypeXP(Protocol):
    @property
    def _np_dtype(self, /) -> _DTypeNP: ...

# https://github.com/pydata/sparse/blob/cc3c8d9/sparse/mlir_backend/_dtypes.py
class _DTypeSparseMlir(Protocol):
    @property
    def np_dtype(self, /) -> _DTypeNP: ...

type DType = _DTypeNP | _DTypeTorch | _DTypeXP

Without numpy present, numpy.dtype needs to be replaced with a compatible Protocol.

It could be useful to further define specialized type aliases for the different categories, e.g. signed/unsigned ints, floats, etc.

Unfortunately this approach isn't waterproof, and isn't guaranteed to work for other libraries. Either way, I believe that this is a better approach than the "harmfully useless" DType = Any approach that's currently used by e.g. scipy.

[jorenham]

NumPy's numeric scalar type signatures (1.24 ~ 2.1)

The only way to structurally type (i.e. using Protocols) the array-api-compatible dtypes of numpy, is through its .type: _SCT property. This moves the problem to the scalar types, i.e. np.bool_ and np.number subtypes (and let's pretend that timedelta64 doesn't exist, like everyone seems to do).

I've analyzed the members, and concluded that it should be possible (if we ignore any potential typechecker bugs _{I'm looking at you, mypy}) to create protocols for several categories of array-api dtypes.

What about all the other array-api libs?

This also covers the array-api-strict dtype, as it is a wrapper around np.dtype.

The torch.dtype can't be expressed in terms of a numpy dtype, but shouldn't be much of a problem to structurally type (I actually already did write the protocol, but I still have to test it 😅).

Distinguishable members of the numeric numpy scalar type kinds

	b	i	u	f	c
__class_getitem__()		X	X	X	X
__complex__()	T	T	T	T	X
__index__()	R	X	X
__round__()		X	X	X	R
as_integer_ratio()				X
bit_count()		X	X
is_integer()		X	X	X

legend:

<blank>	X	R	T
missing	exists	runtime-only (<2.2)	type-check-only (<2.2)

problems

• bool_ (b) has no unique members w.r.t. generic or ndarray when type-checking
• int (i) & uint (u) are indistinguishable
• floating (f) is assignable to complexfloating (c) when type-checking, due to the (incorrect) floating.__complex__ annotation (fixed in numpy>=2.2)

solutions

These could be the protocols:

• Bool (b)
  • item() -> builtins.bool
• Int =: SInt | UInt (iu)
  • bit_count() -> int
• Bool | Int (biu)
  • __index__ (requires numpy >= 2.2 & deprecated)
  • item() -> int
• Float (f)
  • as_integer_ratio() -> tuple[int, int]
• Float64
  • __getnewargs__() -> tuple[float]
• Int | Float (iuf)
  • is_integer()
• Bool | Int | Float (biuf)
  • item() -> float
• Complex (c)
  • __complex__ (requires numpy >= 2.2 & deprecated)
• Complex128 (c16)
  • __getnewargs__() -> tuple[float, float]
• Float64 | Complex128 (f8, c16)
  • __getnewargs__() -> tuple[float] | tuple[float, float]
• Number =: Int | Float | Complex (iufc)
  • __round__ (requires numpy >= 2.2)
  • __class_getitem__ (does that even work in protocols?)
• Bool | Int | Float | Complex (biufc) (like complex_co)
  • item() -> complex

---

TLDR

This involves protocols for each of the following scalar types (ordered by level of abstraction), in terms of np.generic subtypes.

• bool_ (b1), float64 (f8), complex128 (c16)
• floating[Any] (f)
• integer[Any] (iu)
• bool_ | integer[Any] (biu), integer[Any] | floating[Any] (iuf)
• number[Any] (iufc)
• bool_ | number[Any] (biufc) (i.e. the "upper"/"supremum" scalar type)

---

Update

This has been implemented in #170

[jorenham]

superseded by https://github.com/data-apis/array-api-typing