View source on GitHub |
Solves systems of linear eqns A X = RHS
, given Cholesky factorizations.
tf.linalg.cholesky_solve(
chol, rhs, name=None
)
Specifically, returns X
from A X = RHS
, where A = L L^T
, L
is the
chol
arg and RHS
is the rhs
arg.
# Solve 10 separate 2x2 linear systems:
A = ... # shape 10 x 2 x 2
RHS = ... # shape 10 x 2 x 1
chol = tf.linalg.cholesky(A) # shape 10 x 2 x 2
X = tf.linalg.cholesky_solve(chol, RHS) # shape 10 x 2 x 1
# tf.matmul(A, X) ~ RHS
X[3, :, 0] # Solution to the linear system A[3, :, :] x = RHS[3, :, 0]
# Solve five linear systems (K = 5) for every member of the length 10 batch.
A = ... # shape 10 x 2 x 2
RHS = ... # shape 10 x 2 x 5
...
X[3, :, 2] # Solution to the linear system A[3, :, :] x = RHS[3, :, 2]
Returns | |
---|---|
Solution to A x = rhs , shape [..., M, K] .
|