tf.linalg.tridiagonal_matmul

TensorFlow 2 version

View source on GitHub

Multiplies tridiagonal matrix by matrix.

tf.linalg.tridiagonal_matmul(
    diagonals, rhs, diagonals_format='compact', name=None
)

diagonals is representation of 3-diagonal NxN matrix, which depends on diagonals_format.

In matrix format, diagonals must be a tensor of shape [..., M, M], with two inner-most dimensions representing the square tridiagonal matrices. Elements outside of the three diagonals will be ignored.

If sequence format, diagonals is list or tuple of three tensors: [superdiag, maindiag, subdiag], each having shape [..., M]. Last element of superdiag first element of subdiag are ignored.

In compact format the three diagonals are brought together into one tensor of shape [..., 3, M], with last two dimensions containing superdiagonals, diagonals, and subdiagonals, in order. Similarly to sequence format, elements diagonals[..., 0, M-1] and diagonals[..., 2, 0] are ignored.

The sequence format is recommended as the one with the best performance.

rhs is matrix to the right of multiplication. It has shape [..., M, N].

Example:

superdiag = tf.constant([-1, -1, 0], dtype=tf.float64)
maindiag = tf.constant([2, 2, 2], dtype=tf.float64)
subdiag = tf.constant([0, -1, -1], dtype=tf.float64)
diagonals = [superdiag, maindiag, subdiag]
rhs = tf.constant([[1, 1], [1, 1], [1, 1]], dtype=tf.float64)
x = tf.linalg.tridiagonal_matmul(diagonals, rhs, diagonals_format='sequence')