BUG: `__matmul__` should return `NotImplemented` instead of raising a `TypeError`

[lucascolley]

__matmul__ doesn't work with fallback to reflected operators per https://docs.python.org/3/library/constants.html#NotImplemented

In [1]: import scipy, numpy as np, array_api_strict as xp

In [3]: class HelloOp(scipy.sparse.linalg.LinearOperator):
   ...:     def _matvec(self, x):
   ...:         return x
   ...:
   ...:     def __rmatmul__(self, other):
   ...:         print("hello :)")
   ...:

In [4]: H = HelloOp(dtype=None, shape=(2, 2))

In [5]: np.eye(2) @ H
hello :)

In [6]: xp.eye(2) @ H
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[6], line 1
----> 1 xp.eye(2) @ H

File ~/ghq/github.com/scipy/scipy/.pixi/envs/array-api-cpu/lib/python3.13/site-packages/array_api_strict/_array_object.py:832, in Array.__matmul__(self, other)
    828 def __matmul__(self, other: Array, /) -> Array:
    829     """
    830     Performs the operation __matmul__.
    831     """
--> 832     self._check_type_device(other)
    833     # matmul is not defined for scalars, but without this, we may get
    834     # the wrong error message from asarray.
    835     other = self._check_allowed_dtypes(other, "numeric", "__matmul__")

File ~/ghq/github.com/scipy/scipy/.pixi/envs/array-api-cpu/lib/python3.13/site-packages/array_api_strict/_array_object.py:241, in Array._check_type_device(self, other)
    239 # Disallow subclasses of Python scalars, such as np.float64 and np.complex128
    240 elif type(other) not in (bool, int, float, complex):
--> 241     raise TypeError(f"Expected Array or Python scalar; got {type(other)}")

TypeError: Expected Array or Python scalar; got <class '__main__.HelloOp'>

@ev-br happy to accept a patch for this?

[ev-br]

There seems to be a general story about __rop__ in array-api-strict: arithmetic operators do not seem to be accepting custom classes, at all:

In [17]: class HelloOp(scipy.sparse.linalg.LinearOperator):
    ...:      def _matvec(self, x):
    ...:          return x
    ...:      def __rmatmul__(self, other):
    ...:          print("hello :)")
    ...:      def __add__(self, other):
    ...:          print("hello __add__")
    ...:          return NotImplemented
    ...:      def __radd__(self, other):
    ...:          print("hello __radd__")
    ...:          return NotImplemented
    ...: 

In [18]: xp.eye(2) + HelloOp(dtype=None, shape=(2,2))
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[18], line 1
----> 1 xp.eye(2) + HelloOp(dtype=None, shape=(2,2))

... snip ...
--> 1019     self._check_type_device(other)
... snip

TypeError: Expected Array or Python scalar; got <class '__main__.HelloOp'>

In [19]: HelloOp(dtype=None, shape=(2,2))  + xp.eye(2)
hello __add__                    # <<<<<<< it did something
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[19], line 1
----> 1 HelloOp(dtype=None, shape=(2,2))  + xp.eye(2)

.... snip
--> 1019     self._check_type_device(other)
... snip
TypeError: Expected Array or Python scalar; got <class '__main__.HelloOp'>

At the same time, numpy does something else either:

In [20]: np.eye(2) + HelloOp(dtype=None, shape=(2,2))
hello __radd__
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[20], line 1
----> 1 np.eye(2) + HelloOp(dtype=None, shape=(2,2))

TypeError: Concatenation operation is not implemented for NumPy arrays, use np.concatenate() instead. Please do not rely on this error; it may not be given on all Python implementations.

In [21]: HelloOp(dtype=None, shape=(2,2))  + np.eye(2)
hello __add__
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[21], line 1
----> 1 HelloOp(dtype=None, shape=(2,2))  + np.eye(2)

TypeError: operand 'HelloOp' does not support ufuncs (__array_ufunc__=None)

Reading https://data-apis.org/array-api/latest/API_specification/array_object.html#arithmetic-operators I don't see much about custom objects either.

So it seems if we do something about __rmatmul__, we should do it consistently for all __rop__s. And it should probably start from a clarification in the spec.

[lucascolley]

alright, one for the next community meeting I suppose

[ev-br]

TBH, I'd expect it to be a can of fairly ugly worms:

In [17]: np.asarray(3) + xp.asarray(3)
Out[17]: np.int64(6)

In [18]: xp.asarray(3) + np.asarray(3)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[18], line 1
----> 1 xp.asarray(3) + np.asarray(3)

File ~/miniforge3/envs/scipy-dev/lib/python3.12/site-packages/array_api_strict/_array_object.py:542, in Array.__add__(self, other)
    538 def __add__(self, other: Array | complex, /) -> Array:
    539     """
    540     Performs the operation __add__.
    541     """
--> 542     self._check_type_device(other)
    543     other = self._check_allowed_dtypes(other, "numeric", "__add__")
    544     if other is NotImplemented:

File ~/miniforge3/envs/scipy-dev/lib/python3.12/site-packages/array_api_strict/_array_object.py:241, in Array._check_type_device(self, other)
    239 # Disallow subclasses of Python scalars, such as np.float64 and np.complex128
    240 elif type(other) not in (bool, int, float, complex):
--> 241     raise TypeError(f"Expected Array or Python scalar; got {type(other)}")

TypeError: Expected Array or Python scalar; got <class 'numpy.ndarray'>

x-ref #102, so there's also __array__ and __array_function__ and whatever other backends define (__torch_function__?). No idea if there's much to specify beyond "prefer to avoid mixing array types".

[lucascolley]

Yeah I think we need a way of rejecting mixing array types (at least mixing with NumPy) while still returning NotImplemented for unknown types

[lucascolley]

This issue was discussed at https://hackmd.io/zn5bvdZTQIeJmb3RW1B-8g#Reflected-operators-and-array-api-strict-Lucas. The resolution seemed to be that we would update the spec to state that array libraries should return NotImplemented for unknown types (not must since NumPy probably won't), and update array-api-strict to obey the should. We could also make strict either implement this or not based on a flag.