update dpnp.linalg.multi_dot implementation (#1729)

* update dpnp.multi_dot * fix pre-commit * update check limitation calls * address comments * use sycl_queue and usm_type for m and s
IntelPython · Feb 23, 2024 · 6c2036f · 6c2036f
1 parent 4b6650f
commit 6c2036f
Show file tree

Hide file tree

Showing 15 changed files with 584 additions and 216 deletions.
diff --git a/dpnp/dpnp_array.py b/dpnp/dpnp_array.py
@@ -540,7 +540,7 @@ def astype(self, dtype, order="K", casting="unsafe", subok=True, copy=True):
             Array data type casting.
         dtype : dtype
             Target data type.
-        order : {'C', 'F', 'A', 'K'}
+        order : {"C", "F", "A", "K"}, optional
             Row-major (C-style) or column-major (Fortran-style) order.
             When ``order`` is 'A', it uses 'F' if ``a`` is column-major and uses 'C' otherwise.
             And when ``order`` is 'K', it keeps strides as closely as possible.

diff --git a/dpnp/dpnp_iface.py b/dpnp/dpnp_iface.py
@@ -56,6 +56,7 @@
     "array_equal",
     "asnumpy",
     "astype",
+    "check_limitations",
     "check_supported_arrays_type",
     "convert_single_elem_array_to_scalar",
     "default_float_type",
@@ -232,6 +233,58 @@ def astype(x1, dtype, order="K", casting="unsafe", copy=True):
     return dpnp_array._create_from_usm_ndarray(array_obj)
 
 
+def check_limitations(
+    order=None, subok=False, like=None, initial=None, where=True
+):
+    """
+    Checking limitation kwargs for their supported values.
+
+    Parameter `order` is only supported with values ``"C"``, ``"F"``
+    and ``None``.
+    Parameter `subok` is only supported with default value ``False``.
+    Parameter `like` is only supported with default value ``None``.
+    Parameter `initial` is only supported with default value ``None``.
+    Parameter `where` is only supported with default value ``True``.
+
+    Raises
+    ------
+    NotImplementedError
+        If any input kwargs is of unsupported value.
+
+    """
+
+    if order in ("A", "a", "K", "k"):
+        raise NotImplementedError(
+            "Keyword argument `order` is supported only with "
+            f"values ``'C'`` and ``'F'``, but got {order}"
+        )
+    if order not in ("C", "c", "F", "f", None):
+        raise ValueError(
+            "Unrecognized `order` keyword value, expecting "
+            f"``'C'`` or ``'F'``, but got {order}"
+        )
+    if like is not None:
+        raise NotImplementedError(
+            "Keyword argument `like` is supported only with "
+            f"default value ``None``, but got {like}"
+        )
+    if subok is not False:
+        raise NotImplementedError(
+            "Keyword argument `subok` is supported only with "
+            f"default value ``False``, but got {subok}"
+        )
+    if initial is not None:
+        raise NotImplementedError(
+            "Keyword argument `initial` is only supported with "
+            f"default value ``None``, but got {initial}"
+        )
+    if where is not True:
+        raise NotImplementedError(
+            "Keyword argument `where` is supported only with "
+            f"default value ``True``, but got {where}"
+        )
+
+
 def check_supported_arrays_type(*arrays, scalar_type=False, all_scalars=False):
     """
     Return ``True`` if each array has either type of scalar,

diff --git a/dpnp/dpnp_iface_arraycreation.py b/dpnp/dpnp_iface_arraycreation.py
@@ -99,43 +99,6 @@
 ]
 
 
-def _check_limitations(order=None, subok=False, like=None):
-    """
-    Checking limitation kwargs for their supported values.
-
-    Parameter `order` is supported only with values ``C``, ``F`` and ``None``.
-    Parameter `subok` is supported only with default value ``False``.
-    Parameter `like` is supported only with default value ``None``.
-
-    Raises
-    ------
-    NotImplementedError
-        If any input kwargs is of unsupported value.
-
-    """
-
-    if order in ("A", "a", "K", "k"):
-        raise NotImplementedError(
-            "Keyword argument `order` is supported only with "
-            f"values ``'C'`` and ``'F'``, but got {order}"
-        )
-    if order not in ("C", "c", "F", "f", None):
-        raise ValueError(
-            "Unrecognized `order` keyword value, expecting "
-            f"``'C'`` or ``'F'``, but got {order}"
-        )
-    if like is not None:
-        raise NotImplementedError(
-            "Keyword argument `like` is supported only with "
-            f"default value ``None``, but got {like}"
-        )
-    if subok is not False:
-        raise NotImplementedError(
-            "Keyword argument `subok` is supported only with "
-            f"default value ``False``, but got {subok}"
-        )
-
-
 def arange(
     start,
     /,
@@ -223,7 +186,7 @@ def arange(
 
     """
 
-    _check_limitations(like=like)
+    dpnp.check_limitations(like=like)
 
     return dpnp_container.arange(
         start,
@@ -343,7 +306,7 @@ def array(
 
     """
 
-    _check_limitations(subok=subok, like=like)
+    dpnp.check_limitations(subok=subok, like=like)
     if ndmin != 0:
         raise NotImplementedError(
             "Keyword argument `ndmin` is supported only with "
@@ -451,7 +414,7 @@ def asanyarray(
 
     """
 
-    _check_limitations(like=like)
+    dpnp.check_limitations(like=like)
 
     return asarray(
         a,
@@ -548,7 +511,7 @@ def asarray(
 
     """
 
-    _check_limitations(like=like)
+    dpnp.check_limitations(like=like)
 
     return dpnp_container.asarray(
         a,
@@ -654,7 +617,7 @@ def ascontiguousarray(
 
     """
 
-    _check_limitations(like=like)
+    dpnp.check_limitations(like=like)
 
     # at least 1-d array has to be returned
     if dpnp.isscalar(a) or hasattr(a, "ndim") and a.ndim == 0:
@@ -768,7 +731,7 @@ def asfortranarray(
 
     """
 
-    _check_limitations(like=like)
+    dpnp.check_limitations(like=like)
 
     # at least 1-d array has to be returned
     if dpnp.isscalar(a) or hasattr(a, "ndim") and a.ndim == 0:
@@ -867,7 +830,7 @@ def copy(
 
     """
 
-    _check_limitations(subok=subok)
+    dpnp.check_limitations(subok=subok)
 
     if dpnp.is_supported_array_type(a):
         sycl_queue_normalized = dpnp.get_normalized_queue_device(
@@ -1176,7 +1139,7 @@ def empty(
 
     """
 
-    _check_limitations(order=order, like=like)
+    dpnp.check_limitations(order=order, like=like)
     return dpnp_container.empty(
         shape,
         dtype=dtype,
@@ -1276,7 +1239,7 @@ def empty_like(
     """
 
     dpnp.check_supported_arrays_type(a)
-    _check_limitations(order=order, subok=subok)
+    dpnp.check_limitations(order=order, subok=subok)
 
     _shape = a.shape if shape is None else shape
     _dtype = a.dtype if dtype is None else dtype
@@ -1385,7 +1348,7 @@ def eye(
 
     """
 
-    _check_limitations(order=order, like=like)
+    dpnp.check_limitations(order=order, like=like)
 
     return dpnp_container.eye(
         N,
@@ -1485,7 +1448,7 @@ def frombuffer(
 
     """
 
-    _check_limitations(like=like)
+    dpnp.check_limitations(like=like)
     return asarray(
         numpy.frombuffer(buffer, dtype=dtype, count=count, offset=offset),
         device=device,
@@ -1609,7 +1572,7 @@ def fromfile(
 
     """
 
-    _check_limitations(like=like)
+    dpnp.check_limitations(like=like)
     return asarray(
         numpy.fromfile(file, dtype=dtype, count=count, sep=sep, offset=offset),
         device=device,
@@ -1725,7 +1688,7 @@ def fromstring(
 
     """
 
-    _check_limitations(like=like)
+    dpnp.check_limitations(like=like)
     return asarray(
         numpy.fromstring(string, dtype=dtype, count=count, sep=sep),
         device=device,
@@ -1819,7 +1782,7 @@ def full(
 
     """
 
-    _check_limitations(order=order, like=like)
+    dpnp.check_limitations(order=order, like=like)
 
     return dpnp_container.full(
         shape,
@@ -1926,7 +1889,7 @@ def full_like(
     """
 
     dpnp.check_supported_arrays_type(a)
-    _check_limitations(order=order, subok=subok)
+    dpnp.check_limitations(order=order, subok=subok)
 
     _shape = a.shape if shape is None else shape
     _dtype = a.dtype if dtype is None else dtype
@@ -2155,7 +2118,7 @@ def identity(
     if n < 0:
         raise ValueError("negative dimensions are not allowed")
 
-    _check_limitations(like=like)
+    dpnp.check_limitations(like=like)
 
     _dtype = dpnp.default_float_type() if dtype is None else dtype
     return dpnp.eye(
@@ -2759,7 +2722,7 @@ def ones(
 
     """
 
-    _check_limitations(order=order, like=like)
+    dpnp.check_limitations(order=order, like=like)
 
     return dpnp_container.ones(
         shape,
@@ -2861,7 +2824,7 @@ def ones_like(
 
     """
     dpnp.check_supported_arrays_type(a)
-    _check_limitations(order=order, subok=subok)
+    dpnp.check_limitations(order=order, subok=subok)
 
     _shape = a.shape if shape is None else shape
     _dtype = a.dtype if dtype is None else dtype
@@ -3347,7 +3310,7 @@ def zeros(
 
     """
 
-    _check_limitations(order=order, like=like)
+    dpnp.check_limitations(order=order, like=like)
 
     return dpnp_container.zeros(
         shape,
@@ -3450,7 +3413,7 @@ def zeros_like(
     """
 
     dpnp.check_supported_arrays_type(a)
-    _check_limitations(order=order, subok=subok)
+    dpnp.check_limitations(order=order, subok=subok)
 
     _shape = a.shape if shape is None else shape
     _dtype = a.dtype if dtype is None else dtype

diff --git a/dpnp/dpnp_iface_linearalgebra.py b/dpnp/dpnp_iface_linearalgebra.py
@@ -82,7 +82,7 @@ def dot(a, b, out=None):
     b : {dpnp.ndarray, usm_ndarray, scalar}
         Second input array. Both inputs `a` and `b` can not be scalars
         at the same time.
-    out : {dpnp.ndarray, usm_ndarray}, optional
+    out : {None, dpnp.ndarray, usm_ndarray}, optional
         Alternative output array in which to place the result. It must have
         the same shape and data type as the expected output and should be
         C-contiguous. If these conditions are not met, an exception is
@@ -345,11 +345,11 @@ def matmul(
 
     Parameters
     ----------
-    x1 : {dpnp_array, usm_ndarray}
+    x1 : {dpnp.ndarray, usm_ndarray}
         First input array.
-    x2 : {dpnp_array, usm_ndarray}
+    x2 : {dpnp.ndarray, usm_ndarray}
         Second input array.
-    out : {dpnp.ndarray, usm_ndarray}, optional
+    out : {None, dpnp.ndarray, usm_ndarray}, optional
         Alternative output array in which to place the result. It must have
         a shape that matches the signature `(n,k),(k,m)->(n,m)` but the type
         (of the calculated values) will be cast if necessary. Default: ``None``.

diff --git a/dpnp/dpnp_iface_manipulation.py b/dpnp/dpnp_iface_manipulation.py
@@ -1202,14 +1202,13 @@ def ravel(a, order="C"):
     x : {dpnp.ndarray, usm_ndarray}
         Input array. The elements in `a` are read in the order specified by
         order, and packed as a 1-D array.
-    order : {'C', 'F'}, optional
-        The elements of `a` are read using this index order. ``C`` means to
+    order : {"C", "F"}, optional
+        The elements of `a` are read using this index order. ``"C"`` means to
         index the elements in row-major, C-style order, with the last axis
         index changing fastest, back to the first axis index changing slowest.
-        ``F`` means to index the elements in column-major, Fortran-style order,
-        with the first index changing fastest, and the last index changing
-        slowest.
-        By default, ``C`` index order is used.
+        ``"F"`` means to index the elements in column-major, Fortran-style
+        order, with the first index changing fastest, and the last index
+        changing slowest. By default, ``"C"`` index order is used.
 
     Returns
     -------
@@ -1313,15 +1312,15 @@ def reshape(a, /, newshape, order="C", copy=None):
         an integer, then the result will be a 1-D array of that length.
         One shape dimension can be -1. In this case, the value is
         inferred from the length of the array and remaining dimensions.
-    order : {'C', 'F'}, optional
+    order : {"C", "F"}, optional
         Read the elements of `a` using this index order, and place the
-        elements into the reshaped array using this index order. 'C'
+        elements into the reshaped array using this index order. ``"C"``
         means to read / write the elements using C-like index order,
         with the last axis index changing fastest, back to the first
-        axis index changing slowest. 'F' means to read / write the
+        axis index changing slowest. ``"F"`` means to read / write the
         elements using Fortran-like index order, with the first index
         changing fastest, and the last index changing slowest. Note that
-        the 'C' and 'F' options take no account of the memory layout of
+        the ``"C"`` and ``"F"`` options take no account of the memory layout of
         the underlying array, and only refer to the order of indexing.
     copy : bool, optional
         Boolean indicating whether or not to copy the input array.

diff --git a/dpnp/dpnp_iface_mathematical.py b/dpnp/dpnp_iface_mathematical.py
@@ -449,7 +449,7 @@ def clip(a, a_min, a_max, *, out=None, order="K", **kwargs):
     a_min, a_max : {dpnp.ndarray, usm_ndarray, None}
         Minimum and maximum value. If ``None``, clipping is not performed on the corresponding edge.
         Only one of `a_min` and `a_max` may be ``None``. Both are broadcast against `a`.
-    out : {dpnp.ndarray, usm_ndarray}, optional
+    out : {None, dpnp.ndarray, usm_ndarray}, optional
         The results will be placed in this array. It may be the input array for in-place clipping.
         `out` must be of the right shape to hold the output. Its type is preserved.
     order : {"C", "F", "A", "K", None}, optional
@@ -614,8 +614,8 @@ def copysign(
     out : ({None, dpnp.ndarray, usm_ndarray}, optional):
         Output array to populate.
         Array must have the correct shape and the expected data type.
-    order : ({'C', 'F', 'A', 'K'}, optional):
-        Memory layout of the newly output array, if parameter `out` is `None`.
+    order : {"C", "F", "A", "K"}, optional
+        Memory layout of the newly output array, if parameter `out` is ``None``.
         Default: "K".
 
     Returns
@@ -2848,7 +2848,7 @@ def sum(
         data type of `a`, the input array elements are cast to the
         specified data type before computing the sum.
         Default: ``None``.
-    out : {dpnp.ndarray, usm_ndarray}, optional
+    out : {None, dpnp.ndarray, usm_ndarray}, optional
         Alternative output array in which to place the result. It must
         have the same shape as the expected output, but the type of
         the output values will be cast if necessary.