tf.linalg.triangular_solve

TensorFlow 2 version

Solves systems of linear equations with upper or lower triangular matrices by backsubstitution.

tf.linalg.triangular_solve(
    matrix, rhs, lower=True, adjoint=False, name=None
)

matrix is a tensor of shape [..., M, M] whose inner-most 2 dimensions form square matrices. If lower is True then the strictly upper triangular part of each inner-most matrix is assumed to be zero and not accessed. If lower is False then the strictly lower triangular part of each inner-most matrix is assumed to be zero and not accessed. rhs is a tensor of shape [..., M, K].

The output is a tensor of shape [..., M, K]. If adjoint is True then the innermost matrices in output satisfy matrix equations matrix[..., :, :] * output[..., :, :] = rhs[..., :, :]. If adjoint is False then the strictly then the innermost matrices in output satisfy matrix equations adjoint(matrix[..., i, k]) * output[..., k, j] = rhs[..., i, j].

Example:

a = tf.constant([[3,  0,  0,  0],
                 [2,  1,  0,  0],
                 [1,  0,  1,  0],
                 [1,  1,  1,  1]], dtype=tf.float32)

b = tf.constant([[4],
                 [2],
                 [4],
                 [2]], dtype=tf.float32)

x = tf.linalg.triangular_solve(a, b, lower=True)
x
# <tf.Tensor: id=257, shape=(4, 1), dtype=float32, numpy=
# array([[ 1.3333334 ],
#        [-0.66666675],
#        [ 2.6666665 ],
#        [-1.3333331 ]], dtype=float32)>

# in python3 one can use `a@x`
tf.matmul(a, x)
# <tf.Tensor: id=263, shape=(4, 1), dtype=float32, numpy=
# array([[4.       ],
#        [2.       ],
#        [4.       ],
#        [1.9999999]], dtype=float32)>

Numpy Compatibility

Equivalent to scipy.linalg.solve_triangular