def test_diag_block_diag(diag1, diag2):
approx(
B.diag(diag1, diag2),
B.concat2d(
[B.dense(diag1), B.zeros(B.dense(diag2))],
[B.zeros(B.dense(diag2)), B.dense(diag2)],
),
)
assert isinstance(B.diag(diag1, diag2), Diagonal)
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 _attempt_diagonal(rows):
# Check whether the result is diagonal.
# Check that the blocks form a square.
if not all([len(row) == len(rows) for row in rows]):
return None
# Collect the diagonal blocks.
diagonal_blocks = []
for r in range(len(rows)):
for c in range(len(rows[0])):
block_shape = B.shape(rows[r][c])
if r == c:
# Keep track of all the diagonal blocks.
diagonal_blocks.append(rows[r][c])
# All blocks on the diagonal must be diagonal or zero.
if not isinstance(rows[r][c], (Diagonal, Zero)):
return None
# All blocks on the diagonal must be square.
if not block_shape[0] == block_shape[1]:
return None
else:
# All blocks not on the diagonal must be zero.
if not isinstance(rows[r][c], Zero):
return None
return Diagonal(B.concat(*[B.diag(x) for x in diagonal_blocks]))
def _check_matmul_diag(a, b):
for tr_a in [False, True]:
for tr_b in [False, True]:
approx(
B.matmul_diag(a, b, tr_a=tr_a, tr_b=tr_b),
B.diag(B.matmul(B.dense(a), B.dense(b), tr_a=tr_a, tr_b=tr_b)),
)
def _project_pattern(self, x, y, pattern):
# Check whether all data is available.
no_missing = all(pattern)
if no_missing:
# All data is available. Nothing to be done.
u = self.u
else:
# Data is missing. Pick the available entries.
y = B.take(y, pattern, axis=1)
# Ensure that `u` remains a structured matrix.
u = Dense(B.take(self.u, pattern))
# Get number of data points and outputs in this part of the data.
n = B.shape(x)[0]
p = sum(pattern)
# Perform projection.
proj_y_partial = B.matmul(y, B.pinv(u), tr_b=True)
proj_y = B.matmul(proj_y_partial, B.inv(self.s_sqrt), tr_b=True)
# Compute projected noise.
u_square = B.matmul(u, u, tr_a=True)
proj_noise = (
self.noise_obs / B.diag(self.s_sqrt) ** 2 * B.diag(B.pd_inv(u_square))
)
# Convert projected noise to weights.
noises = self.model.noises
weights = noises / (noises + proj_noise)
proj_w = B.ones(B.dtype(weights), n, self.m) * weights[None, :]
# Compute Frobenius norm.
frob = B.sum(y ** 2)
frob = frob - B.sum(proj_y_partial * B.matmul(proj_y_partial, u_square))
# Compute regularising term.
reg = 0.5 * (
n * (p - self.m) * B.log(2 * B.pi * self.noise_obs)
+ frob / self.noise_obs
+ n * B.logdet(B.matmul(u, u, tr_a=True))
+ n * 2 * B.logdet(self.s_sqrt)
)
return x, proj_y, proj_w, reg
def multiply(a: Diagonal, b: AbstractMatrix):
assert_compatible(a, b)
# In the case of broadcasting, `B.diag(b)` will not get the diagonal of the
# broadcasted version of `b`, so we exercise extra caution in that case.
rows, cols = B.shape(b)
if rows == 1 or cols == 1:
b_diag = B.squeeze(B.dense(b))
else:
b_diag = B.diag(b)
return Diagonal(a.diag * b_diag)
def test_diag_block_dense(dense1, dense2):
with AssertDenseWarning(concat_warnings):
res = B.diag(dense1, dense2)
approx(
res,
B.concat2d(
[B.dense(dense1), B.zeros(B.dense(dense1))],
[B.zeros(B.dense(dense2)),
B.dense(dense2)],
),
)
assert isinstance(res, Dense)
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 test_fdd_properties():
p = GP(1, EQ())
# Sample observations.
x = B.linspace(0, 5, 5)
y = p(x, 0.1).sample()
# Compute posterior.
p = p | (p(x, 0.1), y)
fdd = p(B.linspace(0, 5, 10), 0.2)
mean, var = fdd.mean, fdd.var
# Check `var_diag`.
fdd = p(B.linspace(0, 5, 10), 0.2)
approx(fdd.var_diag, B.diag(var))
# Check `mean_var`.
fdd = p(B.linspace(0, 5, 10), 0.2)
approx(fdd.mean_var, (mean, var))
# Check `marginals()`.
fdd = p(B.linspace(0, 5, 10), 0.2)
approx(fdd.marginals(), (B.flatten(mean), B.diag(var)))
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 matmul_diag(a: LowRank, b: LowRank, tr_a=False, tr_b=False):
return B.diag(B.matmul(a, b, tr_a=tr_a, tr_b=tr_b))
def dense(a: Diagonal):
if a.dense is None:
a.dense = B.diag(a.diag)
return a.dense
def diag(a: Kronecker):
return B.kron(B.diag(a.left), B.diag(a.right))
def diag(a: Woodbury):
return B.diag(a.diag) + B.diag(a.lr)
def diag(a: Union[Dense, LowerTriangular, UpperTriangular]):
return B.diag(a.mat)
def __init__(self, mat: AbstractMatrix):
Diagonal.__init__(self, B.diag(mat))
def inverse_transform(x):
chol = B.cholesky(B.reg(x))
return B.concat(B.log(B.diag(chol)), B.tril_to_vec(chol,
offset=-1))
def logdet(a: Union[Diagonal, LowerTriangular, UpperTriangular]):
return B.sum(B.log(B.diag(a)))
def trace(a: AbstractMatrix):
# The implementation of diagonal is optimised, so this should be efficient.
return B.sum(B.diag(a))
def test_dense_diag(diag1):
approx(B.dense(diag1), B.diag(diag1.diag))
_check_cache(diag1)
def test_diag(check_lazy_shapes):
check_function(B.diag, (Tensor(3),))
check_function(B.diag, (Tensor(3, 3),))
# Test rank check for TensorFlow.
with pytest.raises(ValueError):
B.diag(Tensor().tf())
def test_normal_diagonalise(normal1):
approx(
normal1.diagonalise(),
Normal(normal1.mean, B.diag(B.diag(B.dense(normal1.var)))),
)
def multiply(a: UpperTriangular, b: LowerTriangular):
return Diagonal(B.multiply(B.diag(a), B.diag(b)))
def test_conversion_to_diagonal(dense1):
approx(Diagonal(dense1), B.diag(B.diag(dense1)))
def _check_iqf(a, b, c):
res = B.iqf_diag(a, b, c)
# Check correctness.
approx(res, B.diag(B.iqf(a, b, c)))
def test_normal_marginals(normal1):
mean, var = normal1.marginals()
approx(mean, normal1.mean.squeeze())
approx(var, B.diag(normal1.var))
def test_normal_marginal_credible_bounds(normal1):
mean, lower, upper = normal1.marginal_credible_bounds()
approx(mean, normal1.mean.squeeze())
approx(lower, normal1.mean.squeeze() - 1.96 * B.diag(normal1.var)**0.5)
approx(upper, normal1.mean.squeeze() + 1.96 * B.diag(normal1.var)**0.5)
def transform(x):
log_diag = x[:side]
chol = B.vec_to_tril(x[side:], offset=-1) + B.diag(B.exp(log_diag))
return B.matmul(chol, chol, tr_b=True)
