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