MAINT: Let ndarray.__imatmul__ handle inplace matrix multiplication in the array-api

Bas van Beek · Bas van Beek · commit f32cdf69ead3 · 2022-02-25T18:17:16.000+01:00
diff --git a/numpy/array_api/_array_object.py b/numpy/array_api/_array_object.py
@@ -819,23 +819,13 @@ def __imatmul__(self: Array, other: Array, /) -> Array:
         """
         Performs the operation __imatmul__.
         """
-        # Note: NumPy does not implement __imatmul__.
-
         # matmul is not defined for scalars, but without this, we may get
         # the wrong error message from asarray.
         other = self._check_allowed_dtypes(other, "numeric", "__imatmul__")
         if other is NotImplemented:
             return other
-
-        # __imatmul__ can only be allowed when it would not change the shape
-        # of self.
-        other_shape = other.shape
-        if self.shape == () or other_shape == ():
-            raise ValueError("@= requires at least one dimension")
-        if len(other_shape) == 1 or other_shape[-1] != other_shape[-2]:
-            raise ValueError("@= cannot change the shape of the input array")
-        self._array[:] = self._array.__matmul__(other._array)
-        return self
+        res = self._array.__imatmul__(other._array)
+        return self.__class__._new(res)
 
     def __rmatmul__(self: Array, other: Array, /) -> Array:
         """
