|TensorFlow 1 version||View source on GitHub|
Solves tridiagonal systems of equations.
Compat aliases for migration
See Migration guide for more details.
tf.linalg.tridiagonal_solve( diagonals, rhs, diagonals_format='compact', transpose_rhs=False, conjugate_rhs=False, name=None, partial_pivoting=True )
The input can be supplied in various formats:
compact, specified by the
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.
diagonals are supplied as a tuple or list of three
tensors of shapes
[..., N] representing
superdiagonals, diagonals, and subdiagonals, respectively.
N can be either
M; in the latter case, the last element of superdiagonal and the
first element of subdiagonal will be ignored.
compact format the three diagonals are brought together into one tensor
[..., 3, M], with last two dimensions containing superdiagonals,
diagonals, and subdiagonals, in order. Similarly to
diagonals[..., 0, M-1] and
diagonals[..., 2, 0] are ignored.
compact format is recommended as the one with best performance. In case
you need to cast a tensor into a compact format manually, use
An example for a tensor of shape [m, m]:
rhs = tf.constant([...]) matrix = tf.constant([[...]]) m = matrix.shape dummy_idx = [0, 0] # An arbitrary element to use as a dummy indices = [[[i, i + 1] for i in range(m - 1)] + [dummy_idx], # Superdiagonal [[i, i] for i in range(m)], # Diagonal [dummy_idx] + [[i + 1, i] for i in range(m - 1)]] # Subdiagonal diagonals=tf.gather_nd(matrix, indices) x = tf.linalg.tridiagonal_solve(diagonals, rhs)
Regardless of the
rhs is a tensor of shape
[..., M] or
[..., M, K]. The latter allows to simultaneously solve K systems with the
same left-hand sides and K different right-hand sides. If