ML Community Day is November 9! Join us for updates from TensorFlow, JAX, and more Learn more

tf.linalg.LinearOperatorBlockDiag

Combines one or more LinearOperators in to a Block Diagonal matrix.

Inherits From: LinearOperator, Module

Used in the notebooks

Used in the tutorials

This operator combines one or more linear operators [op1,...,opJ], building a new LinearOperator, whose underlying matrix representation has each operator opi on the main diagonal, and zero's elsewhere.

Shape compatibility

If opj acts like a [batch] matrix Aj, then op_combined acts like the [batch] matrix formed by having each matrix Aj on the main diagonal.

Each opj is required to represent a matrix, and hence will have shape batch_shape_j + [M_j, N_j].

If opj has shape batch_shape_j + [M_j, N_j], then the combined operator has shape broadcast_batch_shape + [sum M_j, sum N_j], where broadcast_batch_shape is the mutual broadcast of batch_shape_j, j = 1,...,J, assuming the intermediate batch shapes broadcast.

Arguments to matmul, matvec, solve, and solvevec may either be single Tensors or lists of Tensors that are interpreted as blocks. The jth element of a blockwise list of Tensors must have dimensions that match opj for the given method. If a list of blocks is input, then a list of blocks is returned as well.

When the opj are not guaranteed to be square, this operator's methods might fail due to the combined operator not being square and/or lack of efficient methods.

# Create a 4 x 4 linear operator combined of two 2 x 2 operators.
operator_1 = LinearOperatorFullMatrix([[1., 2.], [3., 4.]])
operator_2 = LinearOperatorFullMatrix([[1., 0.], [0., 1.]])
operator = LinearOperatorBlockDiag([operator_1, operator_2])

operator.to_dense()
==> [[1., 2., 0., 0.],
     [3., 4., 0., 0.],
     [0., 0., 1., 0.],
     [0., 0., 0., 1.]]

operator.shape
==> [4, 4]

operator.log_abs_determinant()
==> scalar Tensor

x1 = ... # Shape [2, 2] Tensor
x2 = ... # Shape [2, 2] Tensor
x = tf.concat([x1, x2], 0)  # Shape [2, 4] Tensor
operator.matmul(x)
==> tf.concat([operator_1.matmul(x1), operator_2.matmul(x2)])

# Create a 5 x 4 linear operator combining three blocks.
operator_1 = LinearOperatorFullMatrix([[1.], [3.]])
operator_2 = LinearOperatorFullMatrix([[1., 6.]])
operator_3 = LinearOperatorFullMatrix([[2.], [7.]])
operator = LinearOperatorBlockDiag([operator_1, operator_2, operator_3])

operator.to_dense()
==> [[1., 0., 0., 0.],
     [3., 0., 0., 0.],
     [0., 1., 6., 0.],
     [0., 0., 0., 2.]]
     [0., 0., 0., 7.]]

operator.shape
==> [5, 4]


# Create a [2, 3] batch of 4 x 4 linear operators.
matrix_44 = tf.random.normal(shape=[2, 3, 4, 4])
operator_44 = LinearOperatorFullMatrix(matrix)

# Create a [1, 3] batch of 5 x 5 linear operators.
matrix_55 = tf.random.normal(shape=[1, 3, 5, 5])
operator_55 = LinearOperatorFullMatrix(matrix_55)

# Combine to create a [2, 3] batch of 9 x 9 operators.
operator_99 = LinearOperatorBlockDiag([operator_44, operator_55])

# Create a shape [2, 3, 9] vector.
x = tf.random.normal(shape=[2, 3, 9])
operator_99.matmul(x)
==> Shape [2, 3, 9] Tensor

# Create a blockwise list of vectors.
x = [tf.random.normal(shape=[2, 3, 4]), tf.random.normal(shape=[2, 3, 5])]
operator_99.matmul(x)
==> [Shape [2, 3, 4] Tensor, Shape [2, 3, 5] Tensor]

Performance

The performance of LinearOperatorBlockDiag on any operation is equal to the sum of the individual operators' operations.