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

tf.linalg.LinearOperatorTridiag

LinearOperator acting like a [batch] square tridiagonal matrix.

Inherits From: LinearOperator, Module

This operator acts like a [batch] square tridiagonal matrix A with shape [B1,...,Bb, N, N] for some b >= 0. The first b indices index a batch member. For every batch index (i1,...,ib), A[i1,...,ib, : :] is an N x M matrix. This matrix A is not materialized, but for purposes of broadcasting this shape will be relevant.

Example usage:

Create a 3 x 3 tridiagonal linear operator.

superdiag = [3., 4., 5.]
diag = [1., -1., 2.]
subdiag = [6., 7., 8]
operator = tf.linalg.LinearOperatorTridiag(
   [superdiag, diag, subdiag],
   diagonals_format='sequence')
operator.to_dense()
<tf.Tensor: shape=(3, 3), dtype=float32, numpy=
array([[ 1.,  3.,  0.],
       [ 7., -1.,  4.],
       [ 0.,  8.,  2.]], dtype=float32)>
operator.shape
TensorShape([3, 3])

Scalar Tensor output.

operator.log_abs_determinant()
<tf.Tensor: shape=(), dtype=float32, numpy=4.3307333>

Create a [2, 3] batch of 4 x 4 linear operators.

diagonals = tf.random.normal(shape=[2, 3, 3, 4])
operator = tf.linalg.LinearOperatorTridiag(
  diagonals,
  diagonals_format='compact')

Create a shape [2, 1, 4, 2] vector. Note that this shape is compatible since the batch dimensions, [2, 1], are broadcast to operator.batch_shape = [2, 3].

y = tf.random.normal(shape=[2, 1, 4, 2])
x = operator.solve(y)
x
<tf.Tensor: shape=(2, 3, 4, 2), dtype=float32, numpy=...,
dtype=float32)>

Shape compatibility

This operator acts on [batch] matrix with compatible shape. x is a batch matrix with compatible shape for matmul and solve if

operator.shape = [B1,...,Bb] + [N, N],  with b >= 0
x.shape =   [C1,...,Cc] + [N, R],
and [C1,...,Cc] broadcasts with [B1,...,Bb].

Performance

Suppose operator is a LinearOperatorTridiag of shape [N, N], and x.shape = [N, R]. Then

  • operator.matmul(x) will take O(N * R) time.
  • operator.solve(x) will take O(N * R) time.

If instead operator and x have shape [B1,...,Bb, N, N] and [B1,...,Bb, N, R], every operation increases in complexity by B1*...*Bb.

Matrix property hints

This LinearOperator is initialized with boolean flags of the form is_X, for X = no