Merge master into improve_perf_det_slogdet

Author: vlad-perevezentsev

CHANGELOG.md

@@ -16,6 +16,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Added implementation of `dpnp.ndarray.data` and `dpnp.ndarray.data.ptr` attributes [#2521](https://github.com/IntelPython/dpnp/pull/2521)
 * Added `dpnp.ndarray.__contains__` method [#2534](https://github.com/IntelPython/dpnp/pull/2534)
 * Added implementation of `dpnp.linalg.lu_factor` (SciPy-compatible) [#2557](https://github.com/IntelPython/dpnp/pull/2557), [#2565](https://github.com/IntelPython/dpnp/pull/2565)
+* Added implementation of `dpnp.piecewise` [#2550](https://github.com/IntelPython/dpnp/pull/2550)
 
 ### Changed
 
@@ -58,6 +59,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Fixed `dpnp.unique` with 1d input array and `axis=0`, `equal_nan=True` keywords passed where the produced result doesn't collapse the NaNs [#2530](https://github.com/IntelPython/dpnp/pull/2530)
 * Resolved issue when `dpnp.ndarray` constructor is called with `dpnp.ndarray.data` as `buffer` keyword [#2533](https://github.com/IntelPython/dpnp/pull/2533)
 * Fixed `dpnp.linalg.cond` to always return a real dtype [#2547](https://github.com/IntelPython/dpnp/pull/2547)
+* Resolved the issue in `dpnp.random` functions to allow any value of `size` where each element is castable to `Py_ssize_t` type [#2578](https://github.com/IntelPython/dpnp/pull/2578)
 
 ### Security
 

dpnp/backend/extensions/lapack/getrs.cpp

@@ -166,6 +166,7 @@ std::pair<sycl::event, sycl::event>
           const dpctl::tensor::usm_ndarray &a_array,
           const dpctl::tensor::usm_ndarray &ipiv_array,
           const dpctl::tensor::usm_ndarray &b_array,
+          oneapi::mkl::transpose trans,
           const std::vector<sycl::event> &depends)
 {
     const int a_array_nd = a_array.get_ndim();
@@ -264,12 +265,6 @@ std::pair<sycl::event, sycl::event>
     const std::int64_t lda = std::max<size_t>(1UL, n);
     const std::int64_t ldb = std::max<size_t>(1UL, n);
 
-    // Use transpose::T if the LU-factorized array is passed as C-contiguous.
-    // For F-contiguous we use transpose::N.
-    oneapi::mkl::transpose trans = is_a_array_c_contig
-                                       ? oneapi::mkl::transpose::T
-                                       : oneapi::mkl::transpose::N;
-
     char *a_array_data = a_array.get_data();
     char *b_array_data = b_array.get_data();
     char *ipiv_array_data = ipiv_array.get_data();

dpnp/backend/extensions/lapack/getrs.hpp

@@ -37,6 +37,7 @@ extern std::pair<sycl::event, sycl::event>
           const dpctl::tensor::usm_ndarray &a_array,
           const dpctl::tensor::usm_ndarray &ipiv_array,
           const dpctl::tensor::usm_ndarray &b_array,
+          oneapi::mkl::transpose trans,
           const std::vector<sycl::event> &depends = {});
 
 extern void init_getrs_dispatch_vector(void);

dpnp/backend/extensions/lapack/lapack_py.cpp

@@ -76,6 +76,13 @@ void init_dispatch_tables(void)
 
 PYBIND11_MODULE(_lapack_impl, m)
 {
+    // Expose oneMKL transpose enum to Python
+    py::enum_<oneapi::mkl::transpose>(m, "Transpose")
+        .value("N", oneapi::mkl::transpose::N)
+        .value("T", oneapi::mkl::transpose::T)
+        .value("C", oneapi::mkl::transpose::C)
+        .export_values(); // Optional, allows access like `Transpose.N`
+
     // Register a custom LinAlgError exception in the dpnp.linalg submodule
     py::module_ linalg_module = py::module_::import("dpnp.linalg");
     py::register_exception<lapack_ext::LinAlgError>(
@@ -160,7 +167,8 @@ PYBIND11_MODULE(_lapack_impl, m)
           "the solves of linear equations with an LU-factored "
           "square coefficient matrix, with multiple right-hand sides",
           py::arg("sycl_queue"), py::arg("a_array"), py::arg("ipiv_array"),
-          py::arg("b_array"), py::arg("depends") = py::list());
+          py::arg("b_array"), py::arg("trans") = oneapi::mkl::transpose::N,
+          py::arg("depends") = py::list());
 
     m.def("_orgqr_batch", &lapack_ext::orgqr_batch,
           "Call `_orgqr_batch` from OneMKL LAPACK library to return "

dpnp/dpnp_iface_functional.py

@@ -36,6 +36,7 @@
 
 """
 
+# pylint: disable=protected-access
 
 from dpctl.tensor._numpy_helper import (
     normalize_axis_index,
@@ -44,7 +45,10 @@
 
 import dpnp
 
-__all__ = ["apply_along_axis", "apply_over_axes"]
+# pylint: disable=no-name-in-module
+from dpnp.dpnp_utils import get_usm_allocations
+
+__all__ = ["apply_along_axis", "apply_over_axes", "piecewise"]
 
 
 def apply_along_axis(func1d, axis, arr, *args, **kwargs):
@@ -266,3 +270,141 @@ def apply_over_axes(func, a, axes):
                 )
         a = res
     return res
+
+
+def piecewise(x, condlist, funclist):
+    """
+    Evaluate a piecewise-defined function.
+
+    Given a set of conditions and corresponding functions, evaluate each
+    function on the input data wherever its condition is true.
+
+    For full documentation refer to :obj:`numpy.piecewise`.
+
+    Parameters
+    ----------
+    x : {dpnp.ndarray, usm_ndarray}
+        The input domain.
+    condlist : {sequence of array-like boolean, bool scalars}
+        Each boolean array/scalar corresponds to a function in `funclist`.
+        Wherever `condlist[i]` is ``True``, `funclist[i](x)` is used as the
+        output value.
+
+        Each boolean array in `condlist` selects a piece of `x`, and should
+        therefore be of the same shape as `x`.
+
+        The length of `condlist` must correspond to that of `funclist`.
+        If one extra function is given, i.e. if
+        ``len(funclist) == len(condlist) + 1``, then that extra function
+        is the default value, used wherever all conditions are ``False``.
+    funclist : {array-like of scalars}
+        A constant value is returned wherever corresponding condition of `x`
+        is ``True``.
+
+    Returns
+    -------
+    out : dpnp.ndarray
+        The output is the same shape and type as `x` and is found by
+        calling the functions in `funclist` on the appropriate portions of `x`,
+        as defined by the boolean arrays in `condlist`. Portions not covered
+        by any condition have a default value of ``0``.
+
+    Limitations
+    -----------
+    Parameters `args` and `kw` are not supported and `funclist` cannot include a
+    callable functions.
+
+    See Also
+    --------
+    :obj:`dpnp.choose` : Construct an array from an index array and a set of
+                         arrays to choose from.
+    :obj:`dpnp.select` : Return an array drawn from elements in `choicelist`,
+                         depending on conditions.
+    :obj:`dpnp.where` : Return elements from one of two arrays depending
+                        on condition.
+
+    Examples
+    --------
+    >>> import dpnp as np
+
+    Define the signum function, which is -1 for ``x < 0`` and +1 for ``x >= 0``.
+
+    >>> x = np.linspace(-2.5, 2.5, 6)
+    >>> np.piecewise(x, [x < 0, x >= 0], [-1, 1])
+    array([-1., -1., -1.,  1.,  1.,  1.])
+
+    """
+    dpnp.check_supported_arrays_type(x)
+    x_dtype = x.dtype
+    if dpnp.is_supported_array_type(condlist) and condlist.ndim in [0, 1]:
+        condlist = [condlist]
+    elif dpnp.isscalar(condlist) or (
+        dpnp.isscalar(condlist[0]) and x.ndim != 0
+    ):
+        # convert scalar to a list of one array
+        # convert list of scalars to a list of one array
+        condlist = [
+            dpnp.full(
+                x.shape, condlist, usm_type=x.usm_type, sycl_queue=x.sycl_queue
+            )
+        ]
+    elif not dpnp.is_supported_array_type(condlist[0]):
+        # convert list of lists to list of arrays
+        # convert list of scalars to a list of 0d arrays (for 0d input)
+        tmp = []
+        for _, cond in enumerate(condlist):
+            tmp.append(
+                dpnp.array(cond, usm_type=x.usm_type, sycl_queue=x.sycl_queue)
+            )
+        condlist = tmp
+
+    dpnp.check_supported_arrays_type(*condlist)
+    if dpnp.is_supported_array_type(funclist):
+        usm_type, exec_q = get_usm_allocations([x, *condlist, funclist])
+    else:
+        usm_type, exec_q = get_usm_allocations([x, *condlist])
+
+    result = dpnp.empty_like(x, usm_type=usm_type, sycl_queue=exec_q)
+
+    condlen = len(condlist)
+    try:
+        if isinstance(funclist, str):
+            raise TypeError("funclist must be a non-string sequence")
+        funclen = len(funclist)
+    except TypeError as e:
+        raise TypeError("funclist must be a sequence of scalars") from e
+
+    if condlen == funclen:
+        # default value is zero
+        default_value = x_dtype.type(0)
+    elif condlen + 1 == funclen:
+        # default value is the last element of funclist
+        default_value = funclist[-1]
+        if callable(default_value):
+            raise NotImplementedError(
+                "Callable functions are not supported currently"
+            )
+        if isinstance(default_value, dpnp.ndarray):
+            default_value = default_value.astype(x_dtype, copy=False)
+        else:
+            default_value = x_dtype.type(default_value)
+        funclist = funclist[:-1]
+    else:
+        raise ValueError(
+            f"with {condlen} condition(s), either {condlen} or {condlen + 1} "
+            "functions are expected"
+        )
+
+    for condition, func in zip(condlist, funclist):
+        if callable(func):
+            raise NotImplementedError(
+                "Callable functions are not supported currently"
+            )
+        if isinstance(func, dpnp.ndarray):
+            func = func.astype(x_dtype, copy=False)
+        else:
+            func = x_dtype.type(func)
+        dpnp.where(condition, func, default_value, out=result)
+        default_value = result
+
+    return result

dpnp/dpnp_utils/dpnp_algo_utils.pyx

@@ -390,10 +390,26 @@ cpdef inline tuple _object_to_tuple(object obj):
     if obj is None:
         return ()
 
-    if cpython.PySequence_Check(obj):
-        return tuple(obj)
+    # dpnp.ndarray unconditionally succeeds in PySequence_Check as it implements __getitem__
+    if cpython.PySequence_Check(obj) and not dpnp.is_supported_array_type(obj):
+        if isinstance(obj, numpy.ndarray):
+            obj = numpy.atleast_1d(obj)
+
+        nd = len(obj)
+        shape = []
+
+        for i in range(0, nd):
+            if cpython.PyBool_Check(obj[i]):
+                raise TypeError("DPNP object_to_tuple(): no item in size can be bool")
+
+            # Assumes each item is castable to Py_ssize_t,
+            # otherwise TypeError will be raised
+            shape.append(<Py_ssize_t> obj[i])
+        return tuple(shape)
 
     if dpnp.isscalar(obj):
+        if cpython.PyBool_Check(obj):
+            raise TypeError("DPNP object_to_tuple(): 'obj' can't be bool")
         return (obj, )
 
     raise ValueError("DPNP object_to_tuple(): 'obj' should be 'None', collections.abc.Sequence, or 'int'")

dpnp/tests/helper.py

@@ -404,6 +404,14 @@ def has_support_aspect64(device=None):
     return dev.has_aspect_fp64
 
 
+def is_arl_or_mtl(device=None):
+    """
+    Return True if a test is running on Arrow Lake or Meteor Lake GPU device,
+    False otherwise.
+    """
+    return _get_dev_mask(device) == 0x7D00
+
+
 def is_cpu_device(device=None):
     """
     Return True if a test is running on CPU device, False otherwise.