Python array API standard typing

Documentation: https://array-api.readthedocs.io

Source Code: https://github.com/34j/types-array-api

Typing for array API and array-api-compat

Installation

Install this via pip (or your favourite package manager):

pip install types-array-api

Usage

Type stubs

Autocompletion for array-api-compat is available in your IDE just by installing this package.

import array_api_compat

xp = array_api_compat.array_namespace(x)

Typing functions using Array

There are multiple ways to type functions:

• from array_api.latest import Array
  
  def simple(x: Array) -> Array:
      return x + 1
  

  The simplest way to enjoy autocompletion for Array. This should be enough for most use cases.

• To make sure that the same type of array is returned (ndarray→ndarray, Tensor→Tensor), a TypeVar bound to Array can be used:

  def generic[TArray: Array](x: TArray) -> TArray:
      return x + 1
  

Advanced Usage

Namespace Type

You can test if an object matches the Protocol as they are runtime-checkable:

import array_api_strict

from array_api.latest import ArrayNamespace, ArrayNamespaceFull


assert isinstance(array_api_strict, ArrayNamespace)
# Full version contains fft and linalg
# fft and linalg are not included by default in array_api_strict
assert not isinstance(array_api_strict, ArrayNamespaceFull)

Shape Typing

• To clarify the input and output shapes, ShapedArray and ShapedAnyArray can be used:

  from array_api.latest import ShapedAnyArray as Array
  
  def sum_last_axis[*TShape](x: Array[*TShape, Any]) -> Array[*TShape]:
      return xp.sum(x, axis=-1)
  

  More complex example using NewType or type aliases:

  RTheta = NewType("RTheta", int)
  XY = NewType("XY", int)
  def polar_coordinates[*TShape](randtheta: Array[*TShape, RTheta]) -> Array[*TShape, XY]:
      """Convert polar coordinates to Cartesian coordinates."""
      r = randtheta[..., 0]
      theta = randtheta[..., 1]
      x = r * xp.cos(theta)
      y = r * xp.sin(theta)
      return xp.stack((x, y), axis=-1)
  

  Note that ShapedAnyArray exists only for documentation purposes and internally it is treated as Array. Using both generic and shaped are impossible due to python/typing#548.

• Note that the below example is ideal but impossible due to Python specification.

  def impossible[
      TDtype,
      TDevice,
      *TShapeFormer: int,
      *TShapeLatter: int,
      TArray: Array
  ](x: TArray[*TShapeFormer, *TShapeLatter | Literal[1], TDtype, TDevice], y: TArray[*TShapeLatter | Literal[1], TDtype, TDevice]) -> TArray[*TShapeFormer, *TShapeLatter, TDtype, TDevice]:
      return x + y # broadcasting
  

Contributors ✨

Thanks goes to these wonderful people (emoji key):

This project follows the all-contributors specification. Contributions of any kind welcome!

Credits

This package was created with Copier and the browniebroke/pypackage-template project template.