Protocol for array objects #229

pmeier · 2021-07-16T07:12:40Z

I've tried to tackle static typing and got a vendorable protocol that can be checked statically as well as at runtime for all but one case I'm going to detail below.

Protocol

import enum
from typing import Any, Optional, Protocol, Tuple, TypeVar, Union, runtime_checkable

A = TypeVar("A")


@runtime_checkable
class VendoredArrayProtocol(Protocol[A]):
    @property
    def dtype(self) -> Any:
        ...

    @property
    def device(self) -> Any:
        ...

    @property
    def ndim(self) -> int:
        ...

    @property
    def shape(self) -> Any:
        ...

    @property
    def size(self) -> int:
        ...

    @property
    def T(self) -> A:
        ...

    def __abs__(self) -> A:
        ...

    def __add__(self, other: Union[int, float, A], /) -> A:
        ...

    def __and__(self, other: Union[bool, int, A], /) -> A:
        ...

    def __array_namespace__(self, /, *, api_version: Optional[str] = None) -> Any:
        ...

    def __bool__(self) -> bool:
        ...

    def __dlpack__(self, /, *, stream: Optional[Union[int, Any]] = None) -> Any:
        ...

    def __dlpack_device__(self) -> Tuple[enum.IntEnum, int]:
        ...

    # This overrides the input type, since object.__eq__ handles any input
    # This overrides the return type, since object.__eq__ returns a bool
    def __eq__(  # type: ignore[override]
        self,
        other: Union[bool, int, float, A],
        /,
    ) -> A:  # type: ignore[override]
        ...

    def __float__(self) -> float:
        ...

    def __floordiv__(self, other: Union[int, float, A], /) -> A:
        ...

    def __ge__(self, other: Union[int, float, A], /) -> A:
        ...

    def __getitem__(
        self,
        key: Union[int, slice, Tuple[Union[int, slice], ...], A],
        /,
    ) -> A:
        ...

    def __gt__(self, other: Union[int, float, A], /) -> A:
        ...

    def __int__(self) -> int:
        ...

    def __invert__(self) -> A:
        ...

    def __le__(self, other: Union[int, float, A], /) -> A:
        ...

    def __len__(self) -> int:
        ...

    def __lshift__(self, other: Union[int, A], /) -> A:
        ...

    def __lt__(self, other: Union[int, float, A], /) -> A:
        ...

    def __matmul__(self, other: A) -> A:
        ...

    def __mod__(self, other: Union[int, float, A], /) -> A:
        ...

    def __mul__(self, other: Union[int, float, A], /) -> A:
        ...

    # This overrides the input type, since object.__ne__ handles any input
    # This overrides the return type, since object.__ne__ returns a bool
    def __ne__(  # type: ignore[override]
        self, other: Union[bool, int, float, A], /
    ) -> A:  # type: ignore[override]
        ...

    def __neg__(self) -> A:
        ...

    def __or__(self, other: Union[bool, int, A], /) -> A:
        ...

    def __pos__(self) -> A:
        ...

    def __pow__(self, other: Union[int, float, A], /) -> A:
        ...

    def __rshift__(self, other: Union[int, A], /) -> A:
        ...

    def __setitem__(
        self,
        key: Union[int, slice, Tuple[Union[int, slice], ...], A],
        value: Union[bool, int, float, A],
        /,
    ) -> None:
        ...

    def __sub__(self, other: Union[int, float, A], /) -> A:
        ...

    def __truediv__(self, other: Union[int, float, A], /) -> A:
        ...

    def __xor__(self, other: Union[bool, int, A], /) -> A:
        ...

To test everything yourself you can use this playground repo.

Current blocker

It is currently impossible to use Ellipsis in type annotations, since its alias ... has a different meaning there. Thus, it is currently impossible to correctly annotate the __getitem__ and __setitem__ methods. There is a fix for this in python/cpython/#22336, but it will only be shipped with Python 3.10. If we leave it out of the annotation, accessing the array with something like Array()[..., 0] will be flagged by mypy although it should be supported according to the specification.

Suggestes improvements

While working on the protocol I found a few issues that could be addressed:

Array.dtype, Array.device, Array.__array_namespace__(), and Array.__dlpack__() should return custom objects, but it is not specified how these objects "look like". In the current state of the protocol I've typed them as Any, but the specification should be more precise.
Array.shape should return Tuple[int, ...], but explicitly check for Python scalars in constant tests array-api-tests#15 (comment) implies that custom objects might also be possible. Maybe we can use Sequence[int]?
The type annotation of the stream parameter from Array.__dlpack__() reads Optional[Union[int, Any]] which is equivalent to Any but more concise.
The binary dunder methods take a specific input types for the other parameter. For example __add__ takes Union[int, float, Array]. IMO they should take Any and return NotImplemented in case they cannot work with the type. For example:
```
class Array:
    def __add__(self, other: Any, /) -> "Array":
        if not isinstance(other, (int, float, Array)):
            return NotImplemented

        # perform addition
```
This makes it harder for static type checkers to catch bugs, because statically something like Array() + None would be allowed, but it gives the other object a chance to work with the Array object by implementing the reflected dunder (here __radd__). If both objects do not know how to deal with the addition, Python will automatically raise a TypeError.

Since the object class defines a __eq__ and __neq__ method according to the proposed scheme above, I needed to put # type: ignore[override] directives in the protocol for the input types.

The text was updated successfully, but these errors were encountered:

BvB93 · 2021-07-16T09:37:59Z

In general this seem like a perfect case for the use of a Protocol, so +1!

Based on my experience in the typing of numpy I do have a couple of comments:

It is currently impossible to use Ellipsis in type annotations, since its alias ... has a different meaning there.

Currently there exists the type-check-only builtins.ellipsis class that is used for this exact purpose.
It's public status is somewhat questionable (python/typeshed#3556), but I'd argue that some pragmatism
is required here, as it is the only way of typing the ... singleton as of the moment.

def __eq__(  # type: ignore[override]
    self,
    other: Union[bool, int, float, A],
    /,
) -> A:  # type: ignore[override]
    ...

While # type: ignore fixes the mypy error, it does not get rid of a more fundamental problem: nearly every single object supports comparison due to object.__eq__. numpy/numpy#17368 (comment) contains a more complete example, but the general gist of it is that we can only describe what happens with VendoredArrayProtocol() == FooType(), but have zero control over FooType() == VendoredArrayProtocol() as this the domain of FooType.__eq__ and/or object.__eq__.

The binary dunder methods take a specific input types for the other parameter. For example __add__ takes Union[int, float, Array]. IMO they should take Any and return NotImplemented in case they cannot work with the type. For example:

As a reminder: for A() + B() to work, the __add__ method of both types does not have to be symmetric. so even if A.__add__ does not support B as argument, everything will work out fine as long as B.__add__ does support A.
A practical example of this in practice would be list.__mul__ and int.__mul__, the former supporting the latter but not the other way around.

Array.shape should return Tuple[int, ...], but explicitly check for Python scalars in constant tests array-api-tests#15 (comment) implies that custom objects might also be possible. Maybe we can use Sequence[int]?

On the subjects of shapes I'd like to point everyone's attention to PEP 646: Variadic Generics. Namelly, it'd allow us to define a single shape type that takes a variable number of parameters (i.e. axis lengths), just like tuple can be Tuple[int], Tuple[int, int], Tuple[int, int, int], etc.. With this mind I'd suggest to either stick to Tuple, or wait until PEP 646 is live and then defined a custom variadic Sequence type or protocol.

Array.dtype, Array.device, Array.__array_namespace__(), and Array.__dlpack__() should return custom objects, but it is not specified how these objects "look like". In the current state of the protocol I've typed them as Any, but the specification should be more precise.

For __dlpack__ specifically it should be quite easy to create a proper return type if we can get builtins.Pycapsule annotated in typeshed (e.g. numpy/numpy#19083 (comment)). If not then it is always possible to create a custom type, but then we would have the issue that it'd have to be distributed from a single central place.

pmeier · 2021-07-16T10:35:42Z

Currently there exists the type-check-only builtins.ellipsis class that is used for this exact purpose.

Awesome, I didn't know or find that. This unblocks this and we could add the protocol to the specification. Current iteration:

Protocol

import enum
from typing import (
    TYPE_CHECKING,
    Any,
    Optional,
    Protocol,
    Tuple,
    TypeVar,
    Union,
    runtime_checkable,
)

if TYPE_CHECKING:
    from builtins import ellipsis


A = TypeVar("A")


@runtime_checkable
class ArrayProtocol(Protocol[A]):
    @property
    def dtype(self) -> Any:
        ...

    @property
    def device(self) -> Any:
        ...

    @property
    def ndim(self) -> int:
        ...

    @property
    def shape(self) -> Tuple[int, ...]:
        ...

    @property
    def size(self) -> int:
        ...

    @property
    def T(self) -> A:
        ...

    def __abs__(self) -> A:
        ...

    def __add__(self, other: Union[int, float, A], /) -> A:
        ...

    def __and__(self, other: Union[bool, int, A], /) -> A:
        ...

    def __array_namespace__(self, /, *, api_version: Optional[str] = None) -> Any:
        ...

    def __bool__(self) -> bool:
        ...

    def __dlpack__(self, /, *, stream: Optional[Union[int, Any]] = None) -> Any:
        ...

    def __dlpack_device__(self) -> Tuple[enum.IntEnum, int]:
        ...

    # This overrides the input type, since object.__eq__ handles any input
    # This overrides the return type, since object.__eq__ returns a bool
    def __eq__(  # type: ignore[override]
        self,
        other: Union[bool, int, float, A],
        /,
    ) -> A:  # type: ignore[override]
        ...

    def __float__(self) -> float:
        ...

    def __floordiv__(self, other: Union[int, float, A], /) -> A:
        ...

    def __ge__(self, other: Union[int, float, A], /) -> A:
        ...

    def __getitem__(
        self,
        key: Union[
            int,
            slice,
            "ellipsis",
            Tuple[Union[int, slice, "ellipsis"], ...],
            A,
        ],
        /,
    ) -> A:
        ...

    def __gt__(self, other: Union[int, float, A], /) -> A:
        ...

    def __int__(self) -> int:
        ...

    def __invert__(self) -> A:
        ...

    def __le__(self, other: Union[int, float, A], /) -> A:
        ...

    def __len__(self) -> int:
        ...

    def __lshift__(self, other: Union[int, A], /) -> A:
        ...

    def __lt__(self, other: Union[int, float, A], /) -> A:
        ...

    def __matmul__(self, other: A) -> A:
        ...

    def __mod__(self, other: Union[int, float, A], /) -> A:
        ...

    def __mul__(self, other: Union[int, float, A], /) -> A:
        ...

    # This overrides the input type, since object.__ne__ handles any input
    # This overrides the return type, since object.__ne__ returns a bool
    def __ne__(  # type: ignore[override]
        self, other: Union[bool, int, float, A], /
    ) -> A:  # type: ignore[override]
        ...

    def __neg__(self) -> A:
        ...

    def __or__(self, other: Union[bool, int, A], /) -> A:
        ...

    def __pos__(self) -> A:
        ...

    def __pow__(self, other: Union[int, float, A], /) -> A:
        ...

    def __rshift__(self, other: Union[int, A], /) -> A:
        ...

    def __setitem__(
        self,
        key: Union[
            int,
            slice,
            "ellipsis",
            Tuple[Union[int, slice, "ellipsis"], ...],
            A,
        ],
        value: Union[bool, int, float, A],
        /,
    ) -> None:
        ...

    def __sub__(self, other: Union[int, float, A], /) -> A:
        ...

    def __truediv__(self, other: Union[int, float, A], /) -> A:
        ...

    def __xor__(self, other: Union[bool, int, A], /) -> A:
        ...

As a reminder: for A() + B() to work, the __add__ method of both types does not have to be symmetric. so even if A.__add__ does not support B as argument, everything will work out fine as long as B.__add__ does support A.

This is only true if B implements __radd__:

class B:
    def __add__(self, other):
        return "B"


class C(B):
    def __radd__(self, other):
        return self.__add__(other)


class A:
    def __add__(self, other):
        if isinstance(other, B):
            return NotImplemented

        return "A"


a = A()
b = B()
c = C()

print(a + c)
print(a + b)

B
TypeError: unsupported operand type(s) for +: 'A' and 'B'

BvB93 · 2021-07-16T11:41:11Z

This is only true if B implements __radd__:

Fair enough, I implicitly assumed that both __add__ and __radd__ would be defined.
Nevertheless, my point here is that there is no need here for Any as long as a custom user type
has properly annotated its __add__ and __radd__ methods (which is the users responsibility).

class A:
    def __add__(self, other: A) -> A: ...
    def __radd__(self, other: A) -> A: ...

class B:
    def __add__(self, other: A | B) -> B: ...
    def __radd__(self, other: A | B) -> B: ...

a: A
b: B

reveal_type(a + a)  # Revealed type is '__main__.A'
reveal_type(a + b)  # Revealed type is '__main__.B'
reveal_type(b + a)  # Revealed type is '__main__.B'
reveal_type(b + b)  # Revealed type is '__main__.B'

asmeurer · 2021-07-16T20:38:38Z

The NotImplemented question came up in the NumPy PR as well (numpy/numpy#18585 (comment)).

My opinion is that the spec should not disallow returning NotImplemented but not require it. Libraries like NumPy have their own dispatching mechanism that isn't based on NotImplemented. I also think it should be out of scope for the spec to define how __eq__ works on unexpected types. An array library may want to work like NumPy where it absorbs things in __eq__, especially if it implements more dtypes than are in the spec, or it may want to return True/False (or NotImplemented), or just error. Any of those options are reasonable. I'm not sure how this affects typing, but it's also my opinion that the type hints should follow from the constraints of the spec, not the other way around.

These are just my opinions though, so maybe others will disagree.

rgommers · 2021-07-27T18:45:37Z

@pmeier can you summarize the issue with if every array library vendors this protocol, instead of putting it in a common package and inheriting from it? That still has a limitation, right?

My opinion is that the spec should not allow returning NotImplemented but not require it.

"not allow" --> "allow" ?

For __dlpack__ specifically it should be quite easy to create a proper return type if we can get builtins.Pycapsule annotated in typeshed (e.g. numpy/numpy#19083 (comment)). If not then it is always possible to create a custom type, but then we would have the issue that it'd have to be distributed from a single central place.

+1 to wait for typeshed here (or go help them along if it starts taking too long).

With this mind I'd suggest to either stick to Tuple, or wait until PEP 646 is live and then defined a custom variadic Sequence type or protocol.

Is there a problem with using Tuple[int, ...]? I'd think that that is a little more specific than plain Tuple, without it being a problem to change it to something even more specific later?

asmeurer · 2021-07-27T23:08:44Z

"not allow" --> "allow" ?

Yes that is what I meant.

pmeier · 2021-07-28T06:01:37Z

@rgommers

can you summarize the issue with if every array library vendors this protocol, instead of putting it in a common package and inheriting from it? That still has a limitation, right?

Not sure what you mean here. If this is about my earlier concern (that I only voiced offline) that typing.TypeVar only works for directly related classes, this is not blocking anymore. Although the documentation states

Alternatively, a type variable may specify an upper bound using bound=<type>. This means that an actual type substituted (explicitly or implicitly) for the type variable must be a subclass of the boundary type [emphasis mine]

it seems to work out fine with the vendored protocol, which by definition cannot be related to an actual array implementation. If you have a look at the playgorund repo and specifically this file, mypy is happy although the ArrayImplementation class is independent of VendoredArrayProtocol.

Is there a problem with using Tuple[int, ...]?

Yes, see data-apis/array-api-tests#15 (comment). But this is not solved by variadic generics either IIUC.

pmeier · 2021-07-28T06:54:55Z

I've tested the protocol against numpy and apart from the device attribute and the array_namespace, dlpack, and dlpack_device dunders, numpy.ndarray's check fine.

BvB93 · 2021-07-28T09:15:36Z

Is there a problem with using Tuple[int, ...]?

Right, to clarify: the reason for excluding subscription was for the sake of creating a compact example.

can you summarize the issue with if every array library vendors this protocol, instead of putting it in a common package and inheriting from it? That still has a limitation, right?

Assuming we won't run into any bugs with type checkers, the sole "issue" here would be code duplication.

rgommers · 2021-07-28T20:52:19Z

Assuming we won't run into any bugs with type checkers, the sole "issue" here would be code duplication.

Good. I'm not bothered by code duplication, I'm pretty sure that all array libraries will prefer that over adding a runtime dependency. However, we still need a way for downstream consumers to access this protocol - downstream packages may be less happy to vendor, and end users even less. That could be solved by:

A standalone package
Each array library exposing its vendored copy

We anyway will need a canonical version somewhere to ensure vendored copies can be updated and remain in sync. So it may be preferable to go with (1).

rgommers · 2021-07-29T18:02:21Z

We had a discussion about this today. A standalone package may indeed be nice, and then we also have a place where other utilities or tests can be added (this has come up several times before, and we always tried to avoid a package because that's work to maintain - but here it seems like we can't avoid that anyway).

In terms of vendoring, it may be possible to avoid that by (in an array library) doing:

try:
   from array_pkgname import ArrayProtocol
except ImportError:
    # this is an optional dependency. static type checking against the Protocol
    # won't work for downstream users if the package is not installed
    class ArrayProtocol():
        pass


class ndarray(ArrayProtocol):
     # array object in this library (could also be named Array, Tensor, etc.)
    ...

A number of libraries don't support static typing at all yet (TensorFlow, CuPy), while some others (NumPy, PyTorch, MXNet) do. Then the only libraries that may want to vendor the protocol are the ones that want to use it in their own test suite.

honno · 2021-12-17T17:59:14Z

Related discussion (somehow totally missed this interesting issue agh)

rgommers added the topic: Static Typing Static typing. label Jul 26, 2021

vnmabus mentioned this issue Aug 26, 2021

ENH: Add an ndarray typing protocol numpy/numpy#19752

Closed

asmeurer mentioned this issue Sep 20, 2021

__array_namespace__ type Hint #267

Open

asmeurer mentioned this issue Mar 18, 2022

ENH: Array API standard and Numpy compatibility numpy/numpy#21135

Open

Illviljan mentioned this issue Aug 9, 2022

Public testing framework for duck array integration pydata/xarray#6894

Open

19 tasks

rgommers mentioned this issue Dec 23, 2022

RFC: add data type inspection utilities to the array API specification #425

Closed

rgommers mentioned this issue Jan 14, 2023

Run-time version of API #584

Open

Illviljan mentioned this issue Jan 19, 2023

Typing of internal datatypes pydata/xarray#7457

Open

rgommers mentioned this issue Apr 12, 2023

use typing.Annotated as standardised format for annotating shapes of tensors using type hints. pytorch/pytorch#98702

Closed

aaronmondal mentioned this issue Sep 15, 2023

Improve ellipsis type #686

Open

asmeurer mentioned this issue Jan 22, 2024

Can we promise to delegate to reverse operators ? #732

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Protocol for array objects #229

Protocol for array objects #229

pmeier commented Jul 16, 2021

BvB93 commented Jul 16, 2021

pmeier commented Jul 16, 2021 •

edited

Loading

BvB93 commented Jul 16, 2021

asmeurer commented Jul 16, 2021 •

edited

Loading

rgommers commented Jul 27, 2021

asmeurer commented Jul 27, 2021

pmeier commented Jul 28, 2021

pmeier commented Jul 28, 2021 •

edited

Loading

BvB93 commented Jul 28, 2021

rgommers commented Jul 28, 2021

rgommers commented Jul 29, 2021

honno commented Dec 17, 2021

Protocol for array objects #229

Protocol for array objects #229

Comments

pmeier commented Jul 16, 2021

Current blocker

Suggestes improvements

BvB93 commented Jul 16, 2021

pmeier commented Jul 16, 2021 • edited Loading

BvB93 commented Jul 16, 2021

asmeurer commented Jul 16, 2021 • edited Loading

rgommers commented Jul 27, 2021

asmeurer commented Jul 27, 2021

pmeier commented Jul 28, 2021

pmeier commented Jul 28, 2021 • edited Loading

BvB93 commented Jul 28, 2021

rgommers commented Jul 28, 2021

rgommers commented Jul 29, 2021

honno commented Dec 17, 2021

pmeier commented Jul 16, 2021 •

edited

Loading

asmeurer commented Jul 16, 2021 •

edited

Loading

pmeier commented Jul 28, 2021 •

edited

Loading