warp.optim.linear.LinearOperator#

class warp.optim.linear.LinearOperator(
shape,
dtype,
device,
matvec,
batch_offsets=None,
max_batch_length=None,
)[source]#

Linear operator to be used as left-hand-side of linear iterative solvers.

Parameters:
  • shape (tuple[int, int]) – Tuple containing the number of rows and columns of the operator

  • dtype (type) – Type of the operator elements

  • device (Device) – Device on which computations involving the operator should be performed

  • matvec (Callable) – Matrix-vector multiplication function

  • batch_offsets (array | None) – Optional array of shape (B+1,) partitioning scalar degrees of freedom into B independent subproblems. batch_offsets[i] is the first scalar degree of freedom of subproblem i. For vector-valued arrays, offsets must be aligned to the vector length. Scalar degrees of freedom at or beyond batch_offsets[-1] are inactive and ignored by batched iterative solvers. When None (default) the operator represents a single subproblem.

  • max_batch_length (int | None) – Optional upper bound on the number of scalar degrees of freedom in any subproblem. Providing this value with batch_offsets allows CUDA reductions to avoid redundant tree levels without accessing the offsets on the host. It must not be smaller than the largest difference between consecutive offsets. Stateful solvers require this value to remain fixed between calls.

The matrix-vector multiplication function should have the following signature:

def matvec(x: warp.array, y: warp.array, z: warp.array, alpha: Scalar, beta: Scalar):
    '''Perform a generalized matrix-vector product.

    This function computes the operation z = alpha * (A @ x) + beta * y, where 'A'
    is the linear operator represented by this class.
    '''
    ...

For performance reasons, by default the iterative linear solvers in this module will try to capture the calls for one or more iterations in CUDA graphs. If the matvec function of a custom LinearOperator cannot be graph-captured, the use_cuda_graph=False parameter should be passed to the solver function.

Attributes

batch_count

Number of independent subproblems.

batch_offsets

Array partitioning the active scalar degrees of freedom, or None.

device

Device on which computations involving the operator are performed.

dtype

Type of the operator elements, such as a scalar type or a matrix block type.

matvec

Matrix-vector multiplication function.

max_batch_length

Upper bound on the number of scalar degrees of freedom in any subproblem, or None.

scalar_type

Scalar component type of the operator's elements.

shape

Number of rows and columns of the operator.

property shape: tuple[int, int][source]#

Number of rows and columns of the operator.

Currently, an operator created by aslinearoperator() from a one-dimensional array of diagonal values retains that array’s one-dimensional shape.

property dtype: type[source]#

Type of the operator elements, such as a scalar type or a matrix block type.

property device: Device[source]#

Device on which computations involving the operator are performed.

property matvec: Callable[source]#

Matrix-vector multiplication function.

Call as operator.matvec(x, y, z, alpha, beta) to store alpha * (A @ x) + beta * y in z, where A is operator.

property scalar_type[source]#

Scalar component type of the operator’s elements.

For scalar elements, scalar_type equals dtype. For vector or matrix elements, scalar_type is the type of each component.

property batch_offsets: array | None[source]#

Array partitioning the active scalar degrees of freedom, or None.

Scalar degrees of freedom at or beyond the final offset are inactive and ignored by batched iterative solvers.

property batch_count: int[source]#

Number of independent subproblems. 1 when batch_offsets is None.

property max_batch_length: int | None[source]#

Upper bound on the number of scalar degrees of freedom in any subproblem, or None.