Make LinAlgError exception to be exposed by LAPACK pybind11 extension

Author: antonwolfy

dpnp/backend/extensions/lapack/lapack_py.cpp

@@ -83,10 +83,9 @@ PYBIND11_MODULE(_lapack_impl, m)
         .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>(
-        linalg_module, "LinAlgError", PyExc_ValueError);
+    // Register a LinAlgError exception in the current submodule
+    py::register_exception<lapack_ext::LinAlgError>(m, "LinAlgError",
+                                                    PyExc_ValueError);
 
     init_dispatch_vectors();
     init_dispatch_tables();

dpnp/linalg/__init__.py

@@ -34,6 +34,9 @@
 """
 
 
+from .dpnp_iface_linalg import (
+    LinAlgError,
+)
 from .dpnp_iface_linalg import __all__ as __all__linalg
 from .dpnp_iface_linalg import (
     cholesky,

dpnp/linalg/dpnp_iface_linalg.py

@@ -38,13 +38,15 @@
 
 # pylint: disable=invalid-name
 # pylint: disable=no-member
+# pylint: disable=no-name-in-module
 
 from typing import NamedTuple
 
 import numpy
 from dpctl.tensor._numpy_helper import normalize_axis_tuple
 
 import dpnp
+from dpnp.backend.extensions.lapack._lapack_impl import LinAlgError
 
 from .dpnp_utils_linalg import (
     assert_2d,
@@ -70,6 +72,7 @@
 )
 
 __all__ = [
+    "LinAlgError",
     "cholesky",
     "cond",
     "cross",
@@ -105,6 +108,9 @@
     "vector_norm",
 ]
 
+# Need to set the module explicitly, since exposed by LAPACK pybind11 extension
+LinAlgError.__module__ = "dpnp.linalg"
+
 
 # pylint:disable=missing-class-docstring
 class EigResult(NamedTuple):
@@ -2330,7 +2336,7 @@ def tensorsolve(a, b, axes=None):
     prod = numpy.prod(old_shape)
 
     if a.size != prod**2:
-        raise dpnp.linalg.LinAlgError(
+        raise LinAlgError(
             "Input arrays must satisfy the requirement "
             "prod(a.shape[b.ndim:]) == prod(a.shape[:b.ndim])"
         )

dpnp/linalg/dpnp_utils_linalg.py

@@ -50,7 +50,6 @@
 import dpnp
 import dpnp.backend.extensions.lapack._lapack_impl as li
 from dpnp.dpnp_utils import get_usm_allocations
-from dpnp.linalg import LinAlgError as LinAlgError
 
 __all__ = [
     "assert_2d",
@@ -943,7 +942,7 @@ def _check_lapack_dev_info(dev_info, error_msg=None):
     if any(dev_info):
         error_msg = error_msg or "Singular matrix"
 
-        raise LinAlgError(error_msg)
+        raise li.LinAlgError(error_msg)
 
 
 def _common_type(*arrays):
@@ -1879,7 +1878,7 @@ def assert_2d(*arrays):
 
     for a in arrays:
         if a.ndim != 2:
-            raise LinAlgError(
+            raise li.LinAlgError(
                 f"{a.ndim}-dimensional array given. The input "
                 "array must be exactly two-dimensional"
             )
@@ -1906,7 +1905,7 @@ def assert_stacked_2d(*arrays):
 
     for a in arrays:
         if a.ndim < 2:
-            raise LinAlgError(
+            raise li.LinAlgError(
                 f"{a.ndim}-dimensional array given. The input "
                 "array must be at least two-dimensional"
             )
@@ -1942,7 +1941,7 @@ def assert_stacked_square(*arrays):
     for a in arrays:
         m, n = a.shape[-2:]
         if m != n:
-            raise LinAlgError(
+            raise li.LinAlgError(
                 "Last 2 dimensions of the input array must be square"
             )
 
@@ -2086,7 +2085,7 @@ def dpnp_cond(x, p=None):
     """Compute the condition number of a matrix."""
 
     if _is_empty_2d(x):
-        raise LinAlgError("cond is not defined on empty arrays")
+        raise li.LinAlgError("cond is not defined on empty arrays")
     if p is None or p == 2 or p == -2:
         s = dpnp.linalg.svd(x, compute_uv=False)
         if p == -2:
@@ -2340,15 +2339,15 @@ def dpnp_lstsq(a, b, rcond=None):
     """
 
     if b.ndim > 2:
-        raise LinAlgError(
+        raise li.LinAlgError(
             f"{b.ndim}-dimensional array given. The input "
             "array must be exactly two-dimensional"
         )
 
     m, n = a.shape[-2:]
     m2 = b.shape[0]
     if m != m2:
-        raise LinAlgError("Incompatible dimensions")
+        raise li.LinAlgError("Incompatible dimensions")
 
     u, s, vh = dpnp_svd(a, full_matrices=False, related_arrays=[b])
 
@@ -2669,20 +2668,20 @@ def dpnp_multi_dot(n, arrays, out=None):
     """Compute dot product of two or more arrays in a single function call."""
 
     if not arrays[0].ndim in [1, 2]:
-        raise LinAlgError(
+        raise li.LinAlgError(
             f"{arrays[0].ndim}-dimensional array given. "
             "First array must be 1-D or 2-D."
         )
 
     if not arrays[-1].ndim in [1, 2]:
-        raise LinAlgError(
+        raise li.LinAlgError(
             f"{arrays[-1].ndim}-dimensional array given. "
             "Last array must be 1-D or 2-D."
         )
 
     for arr in arrays[1:-1]:
         if arr.ndim != 2:
-            raise LinAlgError(
+            raise li.LinAlgError(
                 f"{arr.ndim}-dimensional array given. Inner arrays must be 2-D."
             )