https://github.com/34j/types-array-api
Autogenerated types for array-api-compat and array API - IDE friendly. `pip install types-array-api` and ready to go!
https://github.com/34j/types-array-api
array-api array-api-typing jax mypy mypy-stubs numpy python pytorch type-stubs typing
Last synced: 3 months ago
JSON representation
Autogenerated types for array-api-compat and array API - IDE friendly. `pip install types-array-api` and ready to go!
- Host: GitHub
- URL: https://github.com/34j/types-array-api
- Owner: 34j
- License: apache-2.0
- Created: 2025-06-07T04:11:19.000Z (7 months ago)
- Default Branch: main
- Last Pushed: 2025-10-08T07:00:26.000Z (3 months ago)
- Last Synced: 2025-10-08T08:36:08.181Z (3 months ago)
- Topics: array-api, array-api-typing, jax, mypy, mypy-stubs, numpy, python, pytorch, type-stubs, typing
- Language: Python
- Homepage:
- Size: 946 KB
- Stars: 5
- Watchers: 0
- Forks: 1
- Open Issues: 8
-
Metadata Files:
- Readme: README.md
- Changelog: CHANGELOG.md
- Contributing: CONTRIBUTING.md
- Funding: .github/FUNDING.yml
- License: LICENSE
- Code of conduct: .github/CODE_OF_CONDUCT.md
Awesome Lists containing this project
README
# 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):
```shell
pip install types-array-api
```
## Usage
### Type stubs
Autocompletion for [`array-api-compat`](https://data-apis.org/array-api-compat/) is available in your IDE **just by installing** this package.
```python
import array_api_compat
xp = array_api_compat.array_namespace(x)
```
### Typing functions using `Array`
There are multiple ways to type functions:
- ```python
from array_api._2024_12 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:
```python
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`](https://docs.python.org/3/library/typing.html#typing.runtime_checkable):
```python
import array_api_strict
from array_api._2024_12 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:
```python
from array_api._2024_12 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](https://docs.python.org/3/library/typing.html#newtype) or [type aliases](https://docs.python.org/3/library/typing.html#type-aliases):
```python
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](https://github.com/python/typing/issues/548).
- Note that the below example is ideal but impossible due to Python specification.
```python
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](https://allcontributors.org/docs/en/emoji-key)):
This project follows the [all-contributors](https://github.com/all-contributors/all-contributors) specification. Contributions of any kind welcome!
## Credits
[](https://github.com/copier-org/copier)
This package was created with
[Copier](https://copier.readthedocs.io/) and the
[browniebroke/pypackage-template](https://github.com/browniebroke/pypackage-template)
project template.