ENH: test @ and @= operators in matmul tests

Author: haroune-dev

array_api_tests/test_linalg.py

@@ -400,7 +400,6 @@ def test_inv(x):
 def _test_matmul(namespace, x1, x2):
     matmul = namespace.matmul
 
-    # TODO: Make this also test the @ operator
     if (x1.shape == () or x2.shape == ()
         or len(x1.shape) == len(x2.shape) == 1 and x1.shape != x2.shape
         or len(x1.shape) == 1 and len(x2.shape) >= 2 and x1.shape[0] != x2.shape[-2]
@@ -410,6 +409,8 @@ def _test_matmul(namespace, x1, x2):
         # libraries will use a custom exception class.
         ph.raises(Exception, lambda: xp.matmul(x1, x2),
                "matmul did not raise an exception for invalid shapes")
+        ph.raises(Exception, lambda: x1 @ x2,
+               "@ did not raise an exception for invalid shapes")
         return
     else:
         res = matmul(x1, x2)
@@ -437,6 +438,21 @@ def _test_matmul(namespace, x1, x2):
                                expected=stack_shape + (x1.shape[-2], x2.shape[-1]))
         _test_stacks(matmul, x1, x2, res=res)
 
+    # Test the @ operator against the matmul() result
+    res_op = x1 @ x2
+    ph.assert_dtype("@", in_dtype=[x1.dtype, x2.dtype], out_dtype=res_op.dtype)
+    assert_equal(res, res_op, "@ gives a different result from matmul()")
+
+    # Test @= where the result fits into x1 (same shape and dtype). Only
+    # values are checked: libraries may implement @= as rebinding
+    # (x1 = x1 @ x2) rather than true in-place mutation (numpy itself
+    # does not define __imatmul__), and in-place mutation keeps x1's
+    # dtype while matmul() promotes (e.g. uint8 @= uint16 stays uint8).
+    if res.shape == x1.shape and res.dtype == x1.dtype:
+        x1_inplace = xp.asarray(x1, copy=True)
+        x1_inplace @= x2
+        assert_equal(res, x1_inplace, "@= gives a different result from matmul()")
+
 @pytest.mark.unvectorized
 @pytest.mark.xp_extension('linalg')
 @given(