def triangular_solve(a: UpperTriangular, b: AbstractMatrix, lower_a=True):
if lower_a:
warn_upmodule(
f'Solving against {a}, but "lower_a" is set to "True": ignoring flag.',
category=UserWarning,
)
return B.solve(a, b)
def iqf_diag(a, b, c):
"""Compute the diagonal of `transpose(b) inv(a) c` where `a` is assumed
to be positive definite.
Args:
a (matrix): Matrix `a`.
b (matrix): Matrix `b`.
c (matrix, optional): Matrix `c`. Defaults to `b`.
Returns:
vector: Diagonal of resulting quadratic form.
"""
chol = B.cholesky(a)
chol_b = B.solve(chol, b)
if c is b:
chol_c = chol_b
else:
chol_c = B.solve(chol, c)
return B.matmul_diag(chol_b, chol_c, tr_a=True)
def kl(self, other: "NaturalNormal"):
"""Compute the Kullback-Leibler divergence with respect to another normal
parametrised by its natural parameters.
Args:
other (:class:`.NaturalNormal`): Other.
Returns:
scalar: KL divergence with respect to `other`.
"""
ratio = B.solve(B.chol(self.prec), B.chol(other.prec))
diff = self.mean - other.mean
return 0.5 * (B.sum(ratio**2) - B.logdet(B.mm(
ratio, ratio, tr_a=True)) + B.sum(B.mm(other.prec, diff) * diff) -
B.cast(self.dtype, self.dim))
def iqf_diag(a: Woodbury, b, c):
return B.matmul_diag(b, B.solve(a, c), tr_a=True)
def _check_ratio(a, b):
res = B.ratio(a, b)
# Check correctness.
approx(res, B.trace(B.solve(B.dense(b), B.dense(a))))
def cholesky_solve(a: Union[LowerTriangular, UpperTriangular],
b: AbstractMatrix):
return B.solve(B.transpose(a), B.solve(a, b))
def solve(a: AbstractMatrix, b: AbstractMatrix):
if structured(a, b):
warn_upmodule(f"Solving {a} x = {b}: converting to dense.",
category=ToDenseWarning)
return B.solve(B.dense(a), B.dense(b))
def inverse_transform(x):
eye = B.eye(x)
skew = B.solve(eye + x, eye - x)
return B.tril_to_vec(skew, offset=-1)
def transform(x):
tril = B.vec_to_tril(x, offset=-1)
skew = tril - B.transpose(tril)
eye = B.eye(skew)
return B.solve(eye + skew, eye - skew)
def test_solve_zero(dense_pd, zero2):
assert B.solve(dense_pd, zero2) is zero2
