Merge impl_lu_factor into impl_lu_factor_batch

Author: vlad-perevezentsev

.pre-commit-config.yaml

@@ -8,7 +8,7 @@ repos:
         pass_filenames: false
         args: ["-r", "dpnp", "-lll"]
 -   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v5.0.0
+    rev: v6.0.0
     hooks:
     # Git
     -   id: check-added-large-files

dpnp/backend/extensions/lapack/getrf.cpp

@@ -71,7 +71,7 @@ static sycl::event getrf_impl(sycl::queue &exec_q,
     T *a = reinterpret_cast<T *>(in_a);
 
     const std::int64_t scratchpad_size =
-        mkl_lapack::getrf_scratchpad_size<T>(exec_q, n, n, lda);
+        mkl_lapack::getrf_scratchpad_size<T>(exec_q, m, n, lda);
     T *scratchpad = nullptr;
 
     std::stringstream error_msg;
@@ -88,9 +88,9 @@ static sycl::event getrf_impl(sycl::queue &exec_q,
                   // It must be a non-negative integer.
             n,    // The number of columns in the input matrix A (0 ≤ n).
                   // It must be a non-negative integer.
-            a,    // Pointer to the input matrix A (n x n).
+            a,    // Pointer to the input matrix A (m x n).
             lda,  // The leading dimension of matrix A.
-                  // It must be at least max(1, n).
+                  // It must be at least max(1, m).
             ipiv, // Pointer to the output array of pivot indices.
             scratchpad, // Pointer to scratchpad memory to be used by MKL
                         // routine for storing intermediate results.

dpnp/backend/extensions/lapack/lapack_py.cpp

@@ -135,7 +135,7 @@ PYBIND11_MODULE(_lapack_impl, m)
 
     m.def("_getrf", &lapack_ext::getrf,
           "Call `getrf` from OneMKL LAPACK library to return "
-          "the LU factorization of a general n x n matrix",
+          "the LU factorization of a general m x n matrix",
           py::arg("sycl_queue"), py::arg("a_array"), py::arg("ipiv_array"),
           py::arg("dev_info"), py::arg("depends") = py::list());
 

dpnp/tests/test_linalg.py

@@ -1863,13 +1863,17 @@ class TestLuFactor:
     @staticmethod
     def _apply_pivots_rows(A_dp, piv_dp):
         m = A_dp.shape[0]
-        rows = dpnp.arange(m)
-        for i in range(int(piv_dp.shape[0])):
-            r = int(piv_dp[i].item())
+
+        if m == 0 or piv_dp.size == 0:
+            return A_dp
+
+        rows = list(range(m))
+        piv_np = dpnp.asnumpy(piv_dp)
+        for i, r in enumerate(piv_np):
             if i != r:
-                tmp = rows[i].copy()
-                rows[i] = rows[r]
-                rows[r] = tmp
+                rows[i], rows[r] = rows[r], rows[i]
+
+        rows = dpnp.asarray(rows)
         return A_dp[rows]
 
     @staticmethod
@@ -1955,7 +1959,7 @@ def test_overwrite_copy_special(self):
         a2_orig = a2.copy()
         a2.flags["WRITABLE"] = False
 
-        for a_dp, a_orig in zip((a1, a1), (a1_orig, a2_orig)):
+        for a_dp, a_orig in zip((a1, a2), (a1_orig, a2_orig)):
             lu, piv = dpnp.linalg.lu_factor(
                 a_dp, overwrite_a=True, check_finite=False
             )