def cholesky_solve(chol, rhs, name=None):
"""Solves systems of linear eqns `A X = RHS`, given Cholesky factorizations.
Specifically, returns `X` from `A X = RHS`, where `A = L L^T`, `L` is the
`chol` arg and `RHS` is the `rhs` arg.
```python
# 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]
```
Args:
chol: A `Tensor`. Must be `float32` or `float64`, shape is `[..., M, M]`.
Cholesky factorization of `A`, e.g. `chol = tf.linalg.cholesky(A)`.
For that reason, only the lower triangular parts (including the diagonal)
of the last two dimensions of `chol` are used. The strictly upper part is
assumed to be zero and not accessed.
rhs: A `Tensor`, same type as `chol`, shape is `[..., M, K]`.
name: A name to give this `Op`. Defaults to `cholesky_solve`.
Returns:
Solution to `A x = rhs`, shape `[..., M, K]`.
"""
# To solve C C^* x = rhs, we
# 1. Solve C y = rhs for y, thus y = C^* x
# 2. Solve C^* x = y for x
with ops.name_scope(name, 'cholesky_solve', [chol, rhs]):
y = gen_linalg_ops.matrix_triangular_solve(chol,
rhs,
adjoint=False,
lower=True)
x = gen_linalg_ops.matrix_triangular_solve(chol,
y,
adjoint=True,
lower=True)
return x
def matrix_triangular_solve(matrix, rhs, lower=True, adjoint=False, name=None):
"""Solve systems of linear equations with upper or lower triangular matrices.
`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, N]`.
The output is a tensor of shape `[..., M, N]`. If `adjoint` is `True` then the
innermost matrices in output satisfy matrix equations `
sum_k matrix[..., i, k] * output[..., k, j] = rhs[..., i, j]`.
If `adjoint` is `False` then the
innermost matrices in output satisfy matrix equations
`sum_k 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: shape=(4, 1), dtype=float32, numpy=
array([[ 1.3333334 ],
[-0.66666675],
[ 2.6666665 ],
[-1.3333331 ]], dtype=float32)>
>>> tf.matmul(a, x)
<tf.Tensor: shape=(4, 1), dtype=float32, numpy=
array([[4.],
[2.],
[4.],
[2.]], dtype=float32)>
Args:
matrix: A `Tensor`. Must be one of the following types: `float64`,
`float32`, `half`, `complex64`, `complex128`. Shape is `[..., M, M]`.
rhs: A `Tensor`. Must have the same type as `matrix`. Shape is `[..., M,
N]`.
lower: An optional `bool`. Defaults to `True`. Boolean indicating whether
the innermost matrices in matrix are lower or upper triangular.
adjoint: An optional `bool`. Defaults to `False`. Boolean indicating whether
to solve with matrix or its (block-wise) adjoint.
name: A name for the operation (optional).
Returns:
A `Tensor`. Has the same type as matrix, and shape is `[..., M, N]`.
"""
with ops.name_scope(name, 'triangular_solve', [matrix, rhs]):
return gen_linalg_ops.matrix_triangular_solve(matrix,
rhs,
lower=lower,
adjoint=adjoint)
def target_log_prob(x, y):
counter["target_calls"] += 1
# Corresponds to unnormalized MVN.
# z = matmul(inv(chol(true_cov)), [x, y] - true_mean)
z = array_ops.stack([x, y], axis=-1) - true_mean
z = array_ops.squeeze(gen_linalg_ops.matrix_triangular_solve(
np.linalg.cholesky(true_cov), z[..., array_ops.newaxis]),
axis=-1)
return -0.5 * math_ops.reduce_sum(z**2., axis=-1)
def cholesky_solve(chol, rhs, name=None):
"""Solve linear equations `A X = RHS`, given Cholesky factorization of `A`.
```python
# Solve one system of linear equations (K = 1).
A = [[3, 1], [1, 3]]
RHS = [[2], [22]] # shape 2 x 1
chol = tf.cholesky(A)
X = tf.cholesky_solve(chol, RHS)
# tf.matmul(A, X) ~ RHS
X[:, 0] # Solution to the linear system A x = RHS[:, 0]
# Solve five systems of linear equations (K = 5).
A = [[3, 1], [1, 3]]
RHS = [[1, 2, 3, 4, 5], [11, 22, 33, 44, 55]] # shape 2 x 5
...
X[:, 2] # Solution to the linear system A x = RHS[:, 2]
```
Args:
chol: A `Tensor`. Must be `float32` or `float64`, shape is `[M, M]`.
Cholesky factorization of `A`, e.g. `chol = tf.cholesky(A)`. For that
reason, only the lower triangular part (including the diagonal) of `chol`
is used. The strictly upper part is assumed to be zero and not accessed.
rhs: A `Tensor`, same type as `chol`, shape is `[M, K]`, designating `K`
systems of linear equations.
name: A name to give this `Op`. Defaults to `cholesky_solve`.
Returns:
Solution to `A X = RHS`, shape `[M, K]`. The solutions to the `K` systems.
"""
# To solve C C^* x = rhs, we
# 1. Solve C y = rhs for y, thus y = C^* x
# 2. Solve C^* x = y for x
with ops.op_scope([chol, rhs], name, "cholesky_solve"):
y = gen_linalg_ops.matrix_triangular_solve(chol,
rhs,
adjoint=False,
lower=True)
x = gen_linalg_ops.matrix_triangular_solve(chol,
y,
adjoint=True,
lower=True)
return x
def target_log_prob(x, y):
counter["target_calls"] += 1
# Corresponds to unnormalized MVN.
# z = matmul(inv(chol(true_cov)), [x, y] - true_mean)
z = array_ops.stack([x, y], axis=-1) - true_mean
z = array_ops.squeeze(
gen_linalg_ops.matrix_triangular_solve(
np.linalg.cholesky(true_cov),
z[..., array_ops.newaxis]),
axis=-1)
return -0.5 * math_ops.reduce_sum(z**2., axis=-1)
def cholesky_solve(chol, rhs, name=None):
"""Solves systems of linear eqns `A X = RHS`, given Cholesky factorizations.
```python
# Solve 10 separate 2x2 linear systems:
A = ... # shape 10 x 2 x 2
RHS = ... # shape 10 x 2 x 1
chol = tf.cholesky(A) # shape 10 x 2 x 2
X = tf.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]
```
Args:
chol: A `Tensor`. Must be `float32` or `float64`, shape is `[..., M, M]`.
Cholesky factorization of `A`, e.g. `chol = tf.cholesky(A)`.
For that reason, only the lower triangular parts (including the diagonal)
of the last two dimensions of `chol` are used. The strictly upper part is
assumed to be zero and not accessed.
rhs: A `Tensor`, same type as `chol`, shape is `[..., M, K]`.
name: A name to give this `Op`. Defaults to `cholesky_solve`.
Returns:
Solution to `A x = rhs`, shape `[..., M, K]`.
"""
# To solve C C^* x = rhs, we
# 1. Solve C y = rhs for y, thus y = C^* x
# 2. Solve C^* x = y for x
with ops.name_scope(name, "cholesky_solve", [chol, rhs]):
y = gen_linalg_ops.matrix_triangular_solve(
chol, rhs, adjoint=False, lower=True)
x = gen_linalg_ops.matrix_triangular_solve(
chol, y, adjoint=True, lower=True)
return x
def cholesky_solve(chol, rhs, name=None):
"""Solve linear equations `A X = RHS`, given Cholesky factorization of `A`.
```python
# Solve one system of linear equations (K = 1).
A = [[3, 1], [1, 3]]
RHS = [[2], [22]] # shape 2 x 1
chol = tf.cholesky(A)
X = tf.cholesky_solve(chol, RHS)
# tf.matmul(A, X) ~ RHS
X[:, 0] # Solution to the linear system A x = RHS[:, 0]
# Solve five systems of linear equations (K = 5).
A = [[3, 1], [1, 3]]
RHS = [[1, 2, 3, 4, 5], [11, 22, 33, 44, 55]] # shape 2 x 5
...
X[:, 2] # Solution to the linear system A x = RHS[:, 2]
```
Args:
chol: A `Tensor`. Must be `float32` or `float64`, shape is `[M, M]`.
Cholesky factorization of `A`, e.g. `chol = tf.cholesky(A)`. For that
reason, only the lower triangular part (including the diagonal) of `chol`
is used. The strictly upper part is assumed to be zero and not accessed.
rhs: A `Tensor`, same type as `chol`, shape is `[M, K]`, designating `K`
systems of linear equations.
name: A name to give this `Op`. Defaults to `cholesky_solve`.
Returns:
Solution to `A X = RHS`, shape `[M, K]`. The solutions to the `K` systems.
"""
# To solve C C^* x = rhs, we
# 1. Solve C y = rhs for y, thus y = C^* x
# 2. Solve C^* x = y for x
with ops.name_scope(name, "cholesky_solve", [chol, rhs]):
y = gen_linalg_ops.matrix_triangular_solve(
chol, rhs, adjoint=False, lower=True)
x = gen_linalg_ops.matrix_triangular_solve(
chol, y, adjoint=True, lower=True)
return x
