Type stubs are broken (no numpy support)

Since 1.6 things no longer type check in our code base, since your Array type is incompatible with np.ndarray:

Member	Protocol expects	numpy’s stubs declare
`__array_namespace__`	`api_version: str \| None`, returns `ArrayNamespace[...]`	`api_version: Literal['2021.12', …] \| None`, returns `ModuleType`
`__dlpack__`	`dl_device: tuple[Enum, int] \| None`	`dl_device: tuple[int, int] \| None`
`__dlpack_device__`	returns `tuple[Enum, int]`	returns `tuple[Literal[1], Literal[0]]`

I don’t think this can be fixed other than actually special-casing numpy in the type stubs, unless numpy starts importing your ArrayNamespace and using it as return annotation.

You should run a type checker in CI to see if array_namespace actually typechecks with all the things it supports.

@jorenham

cc @jorenham

Hmm, the Enum doesn't look right, because that's incompatible with int.

Yes, but that matters only in general, not for numpy.

Since numpy will probably never accurately type the return type of __array_namespace__, it needs special handling I think.

/edit: oops clicked wrong.

Is making the __array_namespace__ return type generic an option?

FWIW (and I officially know nothing about type annotations), Enum comes straight from the Array API spec:
https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__dlpack_device__.html

and if anything, -compat should follow the spec. What should be updated, the spec, numpy or -compat layer, I've no idea per the first line above :-).

In [1]: np.array([1,2]).__dlpack_device__()
Out[1]: (1, 0)

For numpy, at least, it's a bare int. So is NumPy wrong, or should that "enum" in the spec be interpreted as a Literal[1,2,...]?

Maybe -compat should type it something like Enum | int?

In [1]: import numpy as np

In [3]: np.ones(3).__dlpack_device__()
Out[3]: (1, 0)

In [4]: import torch

In [5]: torch.ones(3).__dlpack_device__()
Out[5]: (<DLDeviceType.kDLCPU: 1>, 0)

In [6]: import jax.numpy as jnp

In [7]: jnp.ones(3).__dlpack_device__()
An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.
Out[7]: (<DLDeviceType.kDLCPU: 1>, 0)

In [8]: type(jnp.ones(3).__dlpack_device__()[0])
Out[8]: <enum 'DLDeviceType'>

In [9]: type(torch.ones(3).__dlpack_device__()[0])
Out[9]: <enum 'DLDeviceType'>

In [11]: type(cupy.ones(3).__dlpack_device__()[0])
Out[11]: int

-> int would also accept IntEnum (IntEnum subclasses both Enum and int), which is what torch uses: https://github.com/pytorch/pytorch/blob/a97e076079700e9e7e595d59a3e911c37613d8c5/torch/_tensor.py#L1706

PR welcome :-).
Boundary conditions:

-compat needs to support numpy 1.x;
no runtime rewrapping of returns from libraries please;
type annotations are open season.

-compat needs to support numpy 1.x;

Don't you follow SPEC 0 here?

Don't you follow SPEC 0 here?

-compat layer follows whatever active users follow, and some of them are more conservative than SPEC 0 (looking at you scikit-learn).

Ah, luckily the ndarray.__dlpack_device__ return type has been correct since at least 1.22 (e.g.), so -> tuple[int, int] should be fine.

Navigation Menu

Repository navigation

Type stubs are broken (no numpy support) #485

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Type stubs are broken (no numpy support) #485

Description

Activity

lucascolley commented on Oct 26, 2026

jorenham commented on Oct 26, 2026

flying-sheep commented on Oct 26, 2026

jorenham commented on Oct 26, 2026

ev-br commented on Oct 26, 2026

jorenham commented on Oct 26, 2026

ev-br commented on Oct 26, 2026

jorenham commented on Oct 26, 2026

ev-br commented on Oct 26, 2026

jorenham commented on Oct 26, 2026

ev-br commented on Oct 26, 2026

jorenham commented on Oct 26, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions