BlockDiag

class gpjax.linalg.BlockDiag(blocks)

Bases: AbstractLinearOperator

Block diagonal linear operator.

Parameters:

blocks (tuple[AbstractLinearOperator, ...])

as_matrix()

Materialises this linear operator as a matrix.

Note that this can be a computationally (time and/or memory) expensive operation, as many linear operators are defined implicitly, e.g. in terms of their action on a vector.

Arguments: None.

Returns:

A 2-dimensional floating-point JAX array.

in_structure()

Returns the expected input structure of this linear operator.

Arguments: None.

Returns:

A PyTree of jax.ShapeDtypeStruct.

mv(x)

Computes a matrix-vector product between this operator and a vector.

Arguments:

• vector: Should be some PyTree of floating-point arrays, whose structure

  should match self.in_structure().

Returns:

A PyTree of floating-point arrays, with structure that matches self.out_structure().

out_structure()

Returns the expected output structure of this linear operator.

Arguments: None.

Returns:

A PyTree of jax.ShapeDtypeStruct.

transpose()

Transposes this linear operator.

This can be called as either operator.T or operator.transpose().

Arguments: None.

Returns:

Another [lineax.AbstractLinearOperator][].