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 test_dense_lr(lr1):
lr_dense = B.mm(B.dense(lr1.left),
B.dense(lr1.middle),
B.dense(lr1.right),
tr_c=True)
approx(B.dense(lr1), lr_dense)
_check_cache(lr1)
def test_dense_wb(wb1):
lr_dense = B.mm(B.dense(wb1.lr.left),
B.dense(wb1.lr.middle),
B.dense(wb1.lr.right),
tr_c=True)
approx(B.dense(wb1), B.diag(wb1.diag.diag) + lr_dense)
_check_cache(wb1)
def root(a: B.Numeric): # pragma: no cover
"""Compute the positive square root of a positive-definite matrix.
Args:
a (matrix): Matrix to compute square root of.
Returns:
matrix: Positive square root of `a`.
"""
_assert_square_root(a)
u, s, _ = B.svd(a)
return B.mm(u, B.diag(B.sqrt(s)), u, tr_c=True)
def logpdf(self, x):
"""Compute the log-pdf of some data.
Args:
x (column vector): Data to compute log-pdf of.
Returns:
scalar: Log-pdf of `x`.
"""
diff = B.subtract(x, self.mean)
return -0.5 * (-B.logdet(self.prec) + B.cast(self.dtype, self.dim) *
B.cast(self.dtype, B.log_2_pi) +
B.sum(B.mm(self.prec, diff) * diff))
def schur(a: Woodbury):
"""Compute the Schur complement associated to a matrix. A Schur complement will need
to make sense for the type of `a`.
Args:
a (matrix): Matrix to compute Schur complement of.
Returns:
matrix: Schur complement.
"""
if a.schur is None:
second = B.mm(a.lr.right, B.inv(a.diag), a.lr.left, tr_a=True)
a.schur = B.add(B.inv(a.lr.middle), second)
return a.schur
def closest_psd(a, inv=False):
"""Map a matrix to the closest PSD matrix.
Args:
a (tensor): Matrix.
inv (bool, optional): Also invert `a`.
Returns:
tensor: PSD matrix closest to `a` or the inverse of `a`.
"""
a = B.dense(a)
a = (a + B.transpose(a)) / 2
u, s, v = B.svd(a)
signs = B.matmul(u, v, tr_a=True)
s = B.maximum(B.diag(signs) * s, 0)
if inv:
s = B.where(s == 0, 0, 1 / s)
return B.mm(u * B.expand_dims(s, axis=-2), v, tr_b=True)
def iqf(a, b, c):
"""Compute `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:
matrix: 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.mm(chol_b, chol_c, tr_a=True)
def dense(a: LowRank):
if a.dense is None:
a.dense = B.dense(B.mm(a.left, a.middle, a.right, tr_c=True))
return a.dense
def iqf(a: Woodbury, b, c):
return B.mm(b, B.pd_inv(a), c, tr_a=True)
def rand_normal(n=3):
cov = B.randn(n, n)
cov = B.mm(cov, cov, tr_b=True)
return Normal(B.randn(n, 1), cov)