improve coverage

Author: Vahid Tavanashad

dpnp/dpnp_iface_linearalgebra.py

@@ -830,8 +830,7 @@ def matmul(
     >>> np.dot(a, c).shape
     (9, 5, 7, 9, 5, 3)
     >>> np.matmul(a, c).shape
-    (9, 5, 7, 3)
-    >>> # n is 7, k is 4, m is 3
+    (9, 5, 7, 3) # n is 7, k is 4, m is 3
 
     Examples
     --------
@@ -872,7 +871,7 @@ def matmul(
     >>> np.matmul(x1, x2)
     array(-13+0j)
 
-    The ``@`` operator can be used as a shorthand for ``matmul`` on
+    The ``@`` operator can be used as a shorthand for :obj:`dpnp.matmul` on
     :class:`dpnp.ndarray`.
 
     >>> x1 @ x2
@@ -881,6 +880,7 @@ def matmul(
     """
 
     dpnp.check_limitations(subok_linalg=subok, signature=signature)
+    dpnp.check_supported_arrays_type(x1, x2)
 
     return dpnp_multiplication(
         "matmul",
@@ -913,7 +913,7 @@ def matvec(
     Matrix-vector dot product of two arrays.
 
     Given a matrix (or stack of matrices) :math:`\mathbf{A}` in `x1` and
-    a vector (or stack of vectors) :math:`\mathbf{v}` `x2`, the
+    a vector (or stack of vectors) :math:`\mathbf{v}` in `x2`, the
     matrix-vector product is defined as:
 
     .. math::
@@ -952,11 +952,11 @@ def matvec(
 
         Default: ``"K"``.
     axes : {None, list of tuples}, optional
-        A list of tuples with indices of axes the matrix product should operate
-        on. For instance, for the signature of ``(i,j),(j)->(i)``, the base
-        elements are 2d matrices and 1d vectors, where the matrices are assumed
-        to be stored in the last two axes of the first argument, and the
-        vectors in the last axis of the second argument. The corresponding
+        A list of tuples with indices of axes the matrix-vector product should
+        operate on. For instance, for the signature of ``(i,j),(j)->(i)``, the
+        base elements are 2d matrices and 1d vectors, where the matrices are
+        assumed to be stored in the last two axes of the first argument, and
+        the vectors in the last axis of the second argument. The corresponding
         axes keyword would be ``[(-2, -1), (-1,), (-1,)]``.
 
         Default: ``None``.
@@ -1000,6 +1000,7 @@ def matvec(
     """
 
     dpnp.check_limitations(subok_linalg=subok, signature=signature)
+    dpnp.check_supported_arrays_type(x1, x2)
 
     return dpnp_multiplication(
         "matvec",
@@ -1333,7 +1334,7 @@ def vecdot(
 
         Default: ``None``.
     axes : {None, list of tuples}, optional
-        A list of tuples with indices of axes the matrix product should operate
+        A list of tuples with indices of axes the dot product should operate
         on. For instance, for the signature of ``(i),(i)->()``, the base
         elements are vectors and these are taken to be stored in the last axes
         of each argument. The corresponding axes keyword would be
@@ -1414,7 +1415,7 @@ def vecmat(
     Vector-matrix dot product of two arrays.
 
     Given a vector (or stack of vector) :math:`\mathbf{v}` in `x1` and a matrix
-    (or stack of matrices) :math:`\mathbf{A}` `x2`, the vector-matrix product
+    (or stack of matrices) :math:`\mathbf{A}` in `x2`, the vector-matrix product
     is defined as:
 
     .. math::
@@ -1455,12 +1456,12 @@ def vecmat(
 
         Default: ``"K"``.
     axes : {None, list of tuples}, optional
-        A list of tuples with indices of axes the matrix product should operate
-        on. For instance, for the signature of ``(i),(i,j)->(j)``, the base
-        elements are 1D vectors and 2D matrices, where the vectors are assumed
-        to be stored in the last axis of the first argument, and the matrices
-        in the last two axes of the second argument. The corresponding axes
-        keyword would be ``[(-1,), (-2, -1), (-1,)]``.
+        A list of tuples with indices of axes the vector-matrix product should
+        operate on. For instance, for the signature of ``(i),(i,j)->(j)``, the
+        base elements are 1D vectors and 2D matrices, where the vectors are
+        assumed to be stored in the last axis of the first argument, and the
+        matrices in the last two axes of the second argument. The corresponding
+        axes keyword would be ``[(-1,), (-2, -1), (-1,)]``.
 
         Default: ``None``.
 
@@ -1490,14 +1491,18 @@ def vecmat(
     >>> v = np.array([0., 4., 2.])
     >>> a = np.array([[1., 0., 0.],
     ...               [0., 1., 0.],
-    ...               [0., 0., 1.],
-    ...               [0., 6., 8.]])
+    ...               [0., 0., 0.]])
     >>> np.vecmat(v, a)
-    array([ 0.,  4., 0.])
+    array([0., 4., 0.])
 
     """
 
     dpnp.check_limitations(subok_linalg=subok, signature=signature)
+    dpnp.check_supported_arrays_type(x1, x2)
+
+    # cannot directly use dpnp.conj(x1) as it returns `int8` for boolean input
+    if dpnp.issubdtype(x1.dtype, dpnp.complexfloating):
+        x1 = dpnp.conj(x1)
 
     return dpnp_multiplication(
         "vecmat",

dpnp/dpnp_utils/dpnp_utils_linearalgebra.py

@@ -508,11 +508,9 @@ def _gemm_matmul(exec_q, x1, x2, res):
 def _shape_error(shape1, shape2, func, err_msg):
     """Validate the shapes of input and output arrays."""
 
+    # func=None is applicable when err_msg == 2
     if func is not None:
         signature, _ = _get_signature(func)
-    else:
-        # applicable when err_msg == 2
-        assert func is None
 
     if err_msg == 0:
         raise ValueError(
@@ -569,7 +567,13 @@ def _validate_axes(x1, x2, axes, func):
     """Check axes is valid for linear algebra functions."""
 
     def _validate_internal(axes, op, ncores, ndim=None):
-        if ncores == 1:
+        if ncores == 0:
+            if axes != ():
+                raise AxisError(
+                    f"{func}: operand {op} has 0 core dimensions. "
+                    f"Axes item {op} should be an empty tuple."
+                )
+        elif ncores == 1:
             if isinstance(axes, int):
                 axes = (axes,)
             elif not isinstance(axes, tuple):
@@ -616,32 +620,34 @@ def _validate_internal(axes, op, ncores, ndim=None):
     axes[0] = _validate_internal(axes[0], 0, x1_ncore, x1_ndim)
     axes[1] = _validate_internal(axes[1], 1, x2_ncore, x2_ndim)
 
-    if x1_ncore == 1 and x2_ncore == 1:
+    if func == "vecdot":
         if len(axes) == 3:
-            if axes[2] != ():
-                raise AxisError(
-                    f"{func}: output has 0 core dimensions. "
-                    "Axes item 2 should be an empty tuple."
-                )
-        elif len(axes) != 2:
-            raise ValueError(
-                "Axes should be a list of three tuples: two inputs and one "
-                "output. Entry for output can only be omitted if it does not "
-                "have a core axis."
-            )
+            axes[2] = _validate_internal(axes[2], 2, 0)
+            return axes
+
+        if len(axes) == 2:
+            return [axes[0], axes[1], ()]
+
+        raise ValueError(
+            "Axes should be a list of three tuples: two inputs and one "
+            "output. Entry for output can only be omitted if it does not "
+            "have a core axis."
+        )
     else:
         if len(axes) != 3:
             raise ValueError(
                 "Axes should be a list of three tuples: two inputs and one "
                 "output; Entry for output can only be omitted if it does not "
                 "have a core axis."
             )
-        if x1_ncore == 1 or x2_ncore == 1:
+        if x1_ncore == 1 and x2_ncore == 1:
+            axes[2] = _validate_internal(axes[2], 2, 0)
+        elif x1_ncore == 1 or x2_ncore == 1:
             axes[2] = _validate_internal(axes[2], 2, 1)
         else:
             axes[2] = _validate_internal(axes[2], 2, 2)
 
-    return axes
+        return axes
 
 
 def _validate_out_array(out, exec_q):
@@ -829,14 +835,9 @@ def dpnp_multiplication(
 
     """
 
-    dpnp.check_supported_arrays_type(x1, x2)
     res_usm_type, exec_q = get_usm_allocations([x1, x2])
     _validate_out_array(out, exec_q)
 
-    if func == "vecmat":
-        if dpnp.issubdtype(x1.dtype, dpnp.complexfloating):
-            x1 = dpnp.conj(x1)
-
     if order in "aA":
         if x1.flags.fnc and x2.flags.fnc:
             order = "F"
@@ -894,7 +895,7 @@ def dpnp_multiplication(
     )
 
     if axes is not None:
-        # Now that we have result array shape, check axes is within range
+        # Now that result array shape is calculated, check axes is within range
         axes_res = normalize_axis_tuple(axes_res, len(result_shape), "axes")
 
     # Determine the appropriate data types

dpnp/tests/test_product.py

@@ -674,7 +674,7 @@ def setup_method(self):
             ((7, 4, 3), (0, 7, 3, 5)),
         ],
     )
-    def test_matmul(self, order1, order2, shape1, shape2):
+    def test_basic(self, order1, order2, shape1, shape2):
         # input should be float type otherwise they are copied to c-contigous array
         # so testing order becomes meaningless
         dtype = dpnp.default_float_type()
@@ -701,7 +701,7 @@ def test_matmul(self, order1, order2, shape1, shape2):
             "((6, 7, 4, 3), (6, 7, 3, 5))",
         ],
     )
-    def test_matmul_bool(self, shape1, shape2):
+    def test_bool(self, shape1, shape2):
         x = numpy.arange(2, dtype=numpy.bool_)
         a = numpy.resize(x, numpy.prod(shape1)).reshape(shape1)
         b = numpy.resize(x, numpy.prod(shape2)).reshape(shape2)
@@ -766,7 +766,7 @@ def test_axes_1D_ND(self, func, axes):
         expected = getattr(numpy, func)(a, b, axes=axes)
         assert_dtype_allclose(result, expected)
 
-    def test_matmul_axes_1D_1D(self):
+    def test_axes_1D_1D(self):
         a = numpy.arange(3)
         ia = dpnp.array(a)
 
@@ -1417,7 +1417,7 @@ def test_invalid_axes(self, xp):
         assert_raises(AxisError, xp.matmul, a, a, axes=axes)
 
         # axes should be a list of three tuples
-        axes = [0, 0, 0, 0]
+        axes = [0, 0]
         assert_raises(ValueError, xp.matmul, a, a, axes=axes)
 
         a = xp.arange(3 * 4 * 5).reshape(3, 4, 5)
@@ -1940,7 +1940,7 @@ def test_axis2(self):
             [(1,), (1,), ()],
             [(0), (0), ()],
             [0, 1, ()],
-            [-2, -1, ()],
+            [-2, -1],
         ],
     )
     def test_axes(self, axes):
@@ -2145,6 +2145,11 @@ def test_error(self, xp):
         a = xp.ones((5, 5))
         assert_raises(TypeError, xp.vecdot, a, a, axes=[0, 0, ()], axis=-1)
 
+        # axes should be a list of three tuples
+        a = xp.ones(5)
+        axes = [0, 0, 0, 0]
+        assert_raises(ValueError, xp.vecdot, a, a, axes=axes)
+
 
 @testing.with_requires("numpy>=2.2")
 class TestVecmat: