address more comments

Author: vtavana

dpnp/dpnp_utils/dpnp_utils_linearalgebra.py

@@ -58,7 +58,7 @@ def _create_result_array(x1, x2, out, shape, dtype, usm_type, sycl_queue):
             and out.shape == shape
             and out.usm_type == usm_type
             and out.sycl_queue == sycl_queue
-            and (out.flags.c_contiguous or out.flags.f_contiguous)
+            and out.flags.c_contiguous
             and not ti._array_overlap(x1_usm, out_usm)
             and not ti._array_overlap(x2_usm, out_usm)
         ):
@@ -417,7 +417,10 @@ def dpnp_matmul(
         if out is not None:
             dpnp.check_supported_arrays_type(out)
             # out that is passed to the backend should have the correct shape
-            out = dpnp.moveaxis(out, axes_res, (-2, -1))
+            if len(axes_res) == 2:
+                out = dpnp.moveaxis(out, axes_res, (-2, -1))
+            elif len(axes_res) == 1:
+                out = dpnp.moveaxis(out, axes_res, (-1,))
 
     appended_axes = []
     if x1_ndim == 1:
@@ -439,6 +442,37 @@ def dpnp_matmul(
             f"(size {x1_shape[-1]} is different from {x2_shape[-2]})"
         )
 
+    if out is not None:
+        out_shape = out.shape
+        if not appended_axes:
+            if out_shape[-2] != x1_shape[-2]:
+                raise ValueError(
+                    "Output array has a mismatch in its core dimension 0. "
+                    "The core dimensions should follow this signature: (n?,k),(k,m?)->(n?,m?) "
+                    f"(size {out_shape[-2]} is different from {x1_shape[-2]})"
+                )
+            if out_shape[-1] != x2_shape[-1]:
+                raise ValueError(
+                    "Output array has a mismatch in its core dimension 1. "
+                    "The core dimensions should follow this signature: (n?,k),(k,m?)->(n?,m?) "
+                    f"(size {out_shape[-1]} is different from {x2_shape[-1]})"
+                )
+        elif len(appended_axes) == 1:
+            if appended_axes[0] == -1:
+                if out_shape[-1] != x1_shape[-2]:
+                    raise ValueError(
+                        "Output array has a mismatch in its core dimension 0. "
+                        "The core dimensions should follow this signature: (n?,k),(k,m?)->(n?,m?) "
+                        f"(size {out_shape[-1]} is different from {x1_shape[-2]})"
+                    )
+            elif appended_axes[0] == -2:
+                if out_shape[-1] != x2_shape[-1]:
+                    raise ValueError(
+                        "Output array has a mismatch in its core dimension 0. "
+                        "The core dimensions should follow this signature: (n?,k),(k,m?)->(n?,m?) "
+                        f"(size {out_shape[-1]} is different from {x2_shape[-1]})"
+                    )
+
     # Determine the appropriate data types
     gemm_dtype, res_dtype = _op_res_dtype(
         x1, x2, dtype=dtype, casting=casting, sycl_queue=exec_q
@@ -483,10 +517,25 @@ def dpnp_matmul(
                         x2 = dpnp.repeat(x2, x1_shape[i], axis=i)
                 else:
                     raise ValueError(
-                        "arrays could not be broadcast together with remapped shapes."
+                        "Input arrays could not be broadcast together with remapped shapes, "
+                        f"{x1_shape[:-2]} is different from {x2_shape[:-2]}."
                     )
+
         x1_shape = x1.shape
         x2_shape = x2.shape
+        if out is not None:
+            for i in range(x1_ndim - 2):
+                if x1_shape[i] != out_shape[i]:
+                    if not appended_axes:
+                        raise ValueError(
+                            "Output array could not be broadcast together with remapped shapes, "
+                            f"{x1_shape[:-2]} is different from {out_shape[:-2]}."
+                        )
+                    elif len(appended_axes) == 1:
+                        raise ValueError(
+                            "Output array could not be broadcast together with remapped shapes, "
+                            f"{x1_shape[:-2]} is different from {out_shape[:-1]}."
+                        )
         res_shape = tuple(tmp_shape) + (x1_shape[-2], x2_shape[-1])
 
     # handling a special case to provide a similar result to NumPy
@@ -559,6 +608,8 @@ def dpnp_matmul(
 
     if appended_axes:
         result = dpnp.squeeze(result, tuple(appended_axes))
+        if len(appended_axes) == 2 and out is not None:
+            result = dpnp.tile(result, out.shape)
 
     if x1_is_2D and x2_is_2D:
         # add new axes only if one of the input arrays
@@ -572,7 +623,7 @@ def dpnp_matmul(
 
     if out is None:
         if axes is not None:
-            # Move the result to the appropriate axes of out array
+            # Move the data to the appropriate axes of the result array
             if len(axes_res) == 2:
                 result = dpnp.moveaxis(result, (-2, -1), axes_res)
             elif len(axes_res) == 1:
@@ -586,8 +637,7 @@ def dpnp_matmul(
             return result
     else:
         result = dpnp.get_result_array(result, out, casting=casting)
-        if axes is not None:
-            if out is result:
-                # out and out_orig contain the same data but they have different shape
-                return out_orig
+        if axes is not None and out is result:
+            # out and out_orig contain the same data but they have different shape
+            return out_orig
         return result

tests/test_mathematical.py

@@ -2729,6 +2729,29 @@ def test_matmul_axes_out(self, dtype, axes, out_shape):
         # TODO: investigate the effect of factor, see SAT-6700
         assert_dtype_allclose(result, expected, factor=24)
 
+    @pytest.mark.parametrize(
+        "axes, b_shape, out_shape",
+        [
+            ([(1, 0), 0, 0], (3,), (4, 5)),
+            ([(1, 0), 0, 1], (3,), (5, 4)),
+            ([(1, 0), (0, 1), (1, 2)], (3, 1), (5, 4, 1)),
+            ([(1, 0), (0, 1), (0, 2)], (3, 1), (4, 5, 1)),
+            ([(1, 0), (0, 1), (1, 0)], (3, 1), (1, 4, 5)),
+        ],
+    )
+    def test_matmul_axes_out_1D(self, axes, b_shape, out_shape):
+        a = numpy.arange(3 * 4 * 5).reshape(3, 4, 5)
+        b = numpy.arange(3).reshape(b_shape)
+        ia = dpnp.array(a)
+        ib = dpnp.array(b)
+
+        out_dp = dpnp.empty(out_shape)
+        out_np = numpy.empty(out_shape)
+        result = dpnp.matmul(ia, ib, axes=axes, out=out_dp)
+        assert result is out_dp
+        expected = numpy.matmul(a, b, axes=axes, out=out_np)
+        assert_dtype_allclose(result, expected)
+
     @pytest.mark.parametrize("dtype1", get_all_dtypes(no_bool=True))
     @pytest.mark.parametrize(
         "dtype2", get_all_dtypes(no_bool=True, no_none=True)
@@ -2855,6 +2878,25 @@ def test_matmul_out(self, dtype):
         assert result is dpnp_out
         assert_dtype_allclose(result, expected)
 
+    @pytest.mark.parametrize(
+        "out_shape",
+        [
+            ((4, 5)),
+            ((6,)),
+            ((4, 7, 2)),
+        ],
+    )
+    def test_matmul_out_0D(self, out_shape):
+        a = numpy.arange(3)
+        b = dpnp.asarray(a)
+
+        numpy_out = numpy.empty(out_shape)
+        dpnp_out = dpnp.empty(out_shape)
+        result = dpnp.matmul(b, b, out=dpnp_out)
+        expected = numpy.matmul(a, a, out=numpy_out)
+        assert result is dpnp_out
+        assert_dtype_allclose(result, expected)
+
 
 class TestMatmulInvalidCases:
     @pytest.mark.parametrize(
@@ -2892,6 +2934,27 @@ def test_invalid_shape(self, shape_pair):
             with pytest.raises(ValueError):
                 xp.matmul(x1, x2)
 
+    @pytest.mark.parametrize(
+        "shape_pair",
+        [
+            ((5, 4, 3), (3, 1), (3, 4, 1)),
+            ((5, 4, 3), (3, 1), (5, 6, 1)),
+            ((5, 4, 3), (3, 1), (5, 4, 2)),
+            ((5, 4, 3), (3,), (5, 3)),
+            ((5, 4, 3), (3,), (6, 4)),
+            ((3,), (3, 4, 5), (3, 5)),
+            ((3,), (3, 4, 5), (4, 6)),
+        ],
+    )
+    def test_invalid_shape_out(self, shape_pair):
+        for xp in (numpy, dpnp):
+            shape1, shape2, out_shape = shape_pair
+            x1 = xp.arange(numpy.prod(shape1), dtype=xp.float32).reshape(shape1)
+            x2 = xp.arange(numpy.prod(shape2), dtype=xp.float32).reshape(shape2)
+            res = xp.empty(out_shape)
+            with pytest.raises(ValueError):
+                xp.matmul(x1, x2, out=res)
+
     @pytest.mark.parametrize("dtype", get_all_dtypes(no_none=True)[:-2])
     def test_invalid_dtype(self, dtype):
         dpnp_dtype = get_all_dtypes(no_none=True)[-1]

tests/third_party/cupy/math_tests/test_matmul.py

@@ -182,11 +182,11 @@ class TestMatmulLarge(unittest.TestCase):
         rtol=1e-3, atol=1e-3, type_check=has_support_aspect64()
     )  # required for uint8
     def test_operator_matmul(self, xp, dtype1, dtype2):
-        if (dtype1, dtype1) in self.skip_dtypes or (
-            dtype1,
+        if (dtype1, dtype2) in self.skip_dtypes or (
+            dtype2,
             dtype1,
         ) in self.skip_dtypes:
-            return xp.array([])
+            pytest.skip()
         x1 = testing.shaped_random(self.shape_pair[0], xp, dtype1)
         x2 = testing.shaped_random(self.shape_pair[1], xp, dtype2)
         return operator.matmul(x1, x2)
@@ -197,11 +197,11 @@ def test_operator_matmul(self, xp, dtype1, dtype2):
         rtol=1e-3, atol=1e-3, type_check=has_support_aspect64()
     )  # required for uint8
     def test_cupy_matmul(self, xp, dtype1, dtype2):
-        if (dtype1, dtype1) in self.skip_dtypes or (
-            dtype1,
+        if (dtype1, dtype2) in self.skip_dtypes or (
+            dtype2,
             dtype1,
         ) in self.skip_dtypes:
-            return xp.array([])
+            pytest.skip()
         shape1, shape2 = self.shape_pair
         x1 = testing.shaped_random(shape1, xp, dtype1)
         x2 = testing.shaped_random(shape2, xp, dtype2)
@@ -211,7 +211,7 @@ def test_cupy_matmul(self, xp, dtype1, dtype2):
 @pytest.mark.parametrize(
     "shape1, shape2",
     [
-        # TODO: include it when issue #1540 in dpctl is resolved
+        # the first one causes overflow which is undefined behavior
         # ((256, 256, 3, 2), (256, 256, 2, 4)),
         ((256, 256, 3, 2), (2, 4)),
         ((3, 2), (256, 256, 2, 4)),
@@ -233,7 +233,7 @@ def test_cupy_matmul(self, xp, dtype, shape1, shape2):
         return xp.matmul(x1, x2)
 
 
-@pytest.mark.skip("until issue #1540 in dpctl is resolved")
+@pytest.mark.skip("overflow is undefined behavior.")
 class TestMatmulOverflow(unittest.TestCase):
     @testing.for_int_dtypes(name="dtype", no_bool=True)
     @testing.numpy_cupy_allclose(rtol=1e-3, atol=1e-3)  # required for uint8