direct_solver#

nvmath.linalg.direct_solver(
a: AnyTensor,
b: AnyTensor,
*,
options: DirectSolverOptions | dict[str, Any] | None = None,
execution: ExecutionCUDA | str | dict[str, Any] | None = None,
stream: AnyStream | int | None = None,
)[source]#

Solve \(a @ x = b\) for \(x\). This function-form API is a wrapper around the stateful DirectSolver object APIs and is meant for single use (for example, when the user needs to perform just one dense direct linear solve), in which case there is no possibility of amortizing preparatory costs (such as DirectSolver.plan() and DirectSolver.factorize() before DirectSolver.solve()).

Use DirectSolver when reusing a factorization or performing many solves for the same problem size and layout, or when using DirectSolverOptions in a long-lived way (for example a logger or a custom handle).

Parameters:

  • a – The dense operand (or sequence of operands) representing the left-hand side (LHS) of the system of equations: a square matrix a (or batch of square (n, n) matrices) for a @ x = b. The LHS may be a (sequence of) numpy.ndarray, cupy.ndarray, and torch.Tensor. For batched input, a can also be a single tensor of shape (*batch_shape, n, n) (implicit batching) or a list or tuple of matrices of the same shape. Refer to the Semantics section for details.

  • b – The ndarray/tensor or (sequence of ndarray/tensors) representing the dense right-hand side (RHS) of the system of equations. The RHS operand may be a (sequence of) numpy.ndarray, cupy.ndarray, and torch.Tensor. Unbatched shapes are (n,) or (n, nrhs); with batching parallel to a, use a list / tuple of 1D/2D arrays (explicit) or a single tensor (*batch_shape, n, nrhs) (implicit). The package and device must match a. Refer to the Semantics section for details.

  • options – Specify options for the direct linear solver as a DirectSolverOptions object. Alternatively, a dict containing the parameters for the DirectSolverOptions constructor can also be provided. If not specified, the value will be set to the default-constructed DirectSolverOptions object.

  • execution – Specify execution space options for the direct solver as an ExecutionCUDA object. Alternatively, a string ("cuda") or a dict with a "name" of "cuda" and optional device_id relevant to the given execution space. The default execution space is "cuda" and the corresponding ExecutionCUDA will be default-constructed on device 0. CPU operands are executed on the selected CUDA device as needed.

  • stream – Provide the CUDA stream to use for executing the operation. Acceptable inputs include cudaStream_t (as Python int), cupy.cuda.Stream, and torch.cuda.Stream. If a stream is not provided, the current stream from the operand package will be used. See Stream Semantics for more details on stream handling.

Returns:

The result of the specified direct linear solve, x, which has the same batching structure as the RHS b (a single array/tensor, or a sequence of ndarray/tensors, each the same shape as the corresponding tensor in b), a dtype matching a and b, and device and package placement consistent with the chosen execution and operands.

Semantics:

The dense direct linear solver solves \(a @ x = b\) for x given the left-hand side (LHS) a and the right-hand side (RHS) b (square dense \(a\), dense b), using an LU factorization and triangular solves on the target device.

  • Non-batched: a is a square (n, n) matrix; b is a vector (n,) or a matrix (n, nrhs) (one or more right-hand sides). The solve runs on the selected CUDA device using cuSOLVER.

  • Batched (explicit or implicit): the solve uses cuBLAS batched APIs, so in one call every batch entry shares the same system size \(n\) (all (n, n) a). For b, explicit batching uses per-item shapes (n,) or (n, m); implicit batching uses (*batch_shape, n, nrhs). If a and b are implicitly batched, they must share the same leading batch_shape where applicable. The LHS and RHS batching forms may differ (for example, a sequence of a with an implicitly batched b, or the reverse) as long as the number of batch entries is the same for a and b.

See also

nvmath.sparse.advanced.DirectSolver for \(a\) as a sparse CSR matrix (a different feature set and backend).

Note

  • Currently supported dtypes are: float32, float64, complex64, complex128 for both a and b.

  • If users provide options.handle, it is used only for non-batched execution (cuSOLVER path). In batched execution (cuBLAS path), the provided handle is ignored and an internal cached cuBLAS handle is used.

  • When inplace_a=True or inplace_b=True for batched execution, users are responsible for ensuring that the LHS and RHS batch entries do not overlap in memory.

Examples

A single direct_solver() call (equivalent to DirectSolver.plan() followed by DirectSolver.factorize() and DirectSolver.solve() on a stateful object) with NumPy CPU operands, which the implementation executes on the default CUDA device:

>>> import numpy as np
>>> import nvmath.linalg as la
>>>
>>> n = 4
>>> rng = np.random.default_rng(0)
>>> a = rng.standard_normal((n, n), dtype=np.float64) + n * np.eye(n, dtype=np.float64)
>>> b = np.ones(n, dtype=np.float64)
>>>
>>> x = la.direct_solver(a, b)
>>> x.shape
(4,)

For options, stream, and batching, see the Semantics section. More examples (batching, reset_operands(), streams) are in generic direct solver examples on GitHub.

On this page