Skip to content

qttools.kernels.linalg.inv#

[docs] module qttools.kernels.linalg.inv 

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
# Copyright (c) 2024 ETH Zurich and the authors of the qttools package.

from qttools import NDArray, xp
from qttools.profiling import Profiler

profiler = Profiler()


def _invert(a: NDArray) -> NDArray:
    return xp.linalg.inv(a)


def _solve(a: NDArray) -> NDArray:
    return xp.linalg.solve(a, xp.broadcast_to(xp.eye(a.shape[-1]), a.shape))


_inv = _invert
if xp.__name__ == "cupy":
    name = xp.cuda.runtime.getDeviceProperties(0)["name"].decode("utf-8")
    if name.startswith("NVIDIA"):
        from cupy.cublas import set_batched_gesv_limit

        set_batched_gesv_limit(1024)
        _inv = _solve


@profiler.profile(level="debug")
def inv(
    a: NDArray,
) -> NDArray:
    """Computes the (batched) inverse of a matrix.

    Parameters
    ----------
    a : NDArray
        The (batched) matrix.

    Returns
    -------
    NDArray
        The inverse (batched) of the matrix.

    """
    return _inv(a)