support parameter out for matmul func (#722)

* add out to matmul function

Co-authored-by: Sergey Shalnov <[email protected]>

Author: KsanaKozlova

dpnp/dpnp_algo/dpnp_algo.pyx

@@ -174,7 +174,8 @@ cpdef dparray dpnp_init_val(shape, dtype, value):
     return result
 
 
-cpdef dparray dpnp_matmul(dparray in_array1, dparray in_array2):
+cpdef dparray dpnp_matmul(dparray in_array1, dparray in_array2, dparray out=None):
+
     cdef vector[Py_ssize_t] shape_result
 
     cdef vector[Py_ssize_t] shape1 = in_array1.shape
@@ -235,8 +236,18 @@ cpdef dparray dpnp_matmul(dparray in_array1, dparray in_array2):
     cdef DPNPFuncData kernel_data = get_dpnp_function_ptr(DPNP_FN_MATMUL, param1_type, param2_type)
 
     result_type = dpnp_DPNPFuncType_to_dtype( < size_t > kernel_data.return_type)
-    # ceate result array with type given by FPTR data
-    cdef dparray result = dparray(shape_result, dtype=result_type)
+
+    cdef dparray result
+
+    if out is not None:
+        if out.dtype != result_type:
+            checker_throw_value_error('matmul', 'out.dtype', out.dtype, result_type)
+        if out.shape != shape_result:
+            checker_throw_value_error('matmul', 'out.shape', out.shape, shape_result)
+        result = out
+    else:
+        result = dparray(shape_result, dtype=result_type)
+
     if result.size == 0:
         return result
 

dpnp/dpnp_iface.py

@@ -144,7 +144,7 @@ def get_include():
     return dpnp_path
 
 
-def matmul(in_array1, in_array2, out=None):
+def matmul(in_array1, in_array2, out=None, **kwargs):
     """
     Matrix product of two arrays.
 
@@ -154,7 +154,7 @@ def matmul(in_array1, in_array2, out=None):
     -----------
     Input arrays are supported as :obj:`dpnp.ndarray`.
     Otherwise the function will be executed sequentially on CPU.
-    Parameter ``out`` is supported only with default value ``None``.
+    Parameter ``out`` is supported as :obj:`dpnp.ndarray` and as default value ``None``.
     Input array data types are limited by supported DPNP :ref:`Data types`.
 
     See Also
@@ -180,33 +180,29 @@ def matmul(in_array1, in_array2, out=None):
 
     """
 
-    is_dparray1 = isinstance(in_array1, dparray)
-    is_dparray2 = isinstance(in_array2, dparray)
-
-    if (not use_origin_backend(in_array1) and is_dparray1 and is_dparray2):
-
-        if out is not None:
-            checker_throw_value_error("matmul", "out", type(out), None)
-
-        """
-        Cost model checks
-        """
-        cost_size = 4096  # 2D array shape(64, 64)
-        if ((in_array1.dtype == numpy.float64) or (in_array1.dtype == numpy.float32)):
+    if not use_origin_backend(in_array1) and not kwargs:
+        if not isinstance(in_array1, dparray):
+            pass
+        elif not isinstance(in_array2, dparray):
+            pass
+        elif out is not None and not isinstance(out, dparray):
+            pass
+        else:
             """
-            Floating point types are handled via original math library better than SYCL math library
+            Cost model checks
             """
-            cost_size = 262144  # 2D array shape(512, 512)
 
-        dparray1_size = in_array1.size
-        dparray2_size = in_array2.size
+            dparray1_size = in_array1.size
+            dparray2_size = in_array2.size
+            cost_size = 4096  # 2D array shape(64, 64)
 
-        if (dparray1_size > cost_size) and (dparray2_size > cost_size):
-            # print(f"dparray1_size={dparray1_size}")
-            return dpnp_matmul(in_array1, in_array2)
+            if ((in_array1.dtype == numpy.float64) or (in_array1.dtype == numpy.float32)):
+                """
+                Floating point types are handled via original math library better than SYCL math library
+                """
+                cost_size = 262144  # 2D array shape(512, 512)
 
-    input1 = asnumpy(in_array1) if is_dparray1 else in_array1
-    input2 = asnumpy(in_array2) if is_dparray2 else in_array2
+            if (dparray1_size > cost_size) and (dparray2_size > cost_size):
+                return dpnp_matmul(in_array1, in_array2, out=out)
 
-    # TODO need to return dparray instead ndarray
-    return numpy.matmul(input1, input2, out=out)
+    return call_origin(numpy.matmul, in_array1, in_array2, out=out, **kwargs)

tests/test_mixins.py

@@ -46,6 +46,24 @@ def test_matmul2(self):
 
         numpy.testing.assert_array_equal(expected, result)
 
+    def test_matmul3(self):
+        array_data1 = numpy.full((513, 513), 5)
+        array_data2 = numpy.full((513, 513), 2)
+        out = numpy.empty((513, 513), dtype=numpy.float64)
+
+        # DPNP
+        array1 = inp.array(array_data1, dtype=inp.float64)
+        array2 = inp.array(array_data2, dtype=inp.float64)
+        out1 = inp.array(out, dtype=inp.float64)
+        result = inp.matmul(array1, array2, out=out1)
+
+        # original
+        array_1 = numpy.array(array_data1, dtype=numpy.float64)
+        array_2 = numpy.array(array_data2, dtype=numpy.float64)
+        expected = numpy.matmul(array_1, array_2, out=out)
+
+        numpy.testing.assert_array_equal(expected, result)
+
 
 if __name__ == '__main__':
     unittest.main()