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 test_normal_arithmetic():
chol = np.random.randn(3, 3)
dist = Normal(chol.dot(chol.T), np.random.randn(3, 1))
chol = np.random.randn(3, 3)
dist2 = Normal(chol.dot(chol.T), np.random.randn(3, 1))
A = np.random.randn(3, 3)
a = np.random.randn(1, 3)
b = 5.
# Test matrix multiplication.
allclose((dist.rmatmul(a)).mean, dist.mean.dot(a))
allclose((dist.rmatmul(a)).var, a.dot(B.dense(dist.var)).dot(a.T))
allclose((dist.lmatmul(A)).mean, A.dot(dist.mean))
allclose((dist.lmatmul(A)).var, A.dot(B.dense(dist.var)).dot(A.T))
# Test multiplication.
allclose((dist * b).mean, dist.mean * b)
allclose((dist * b).var, dist.var * b**2)
allclose((b * dist).mean, dist.mean * b)
allclose((b * dist).var, dist.var * b**2)
with pytest.raises(NotImplementedError):
dist.__mul__(dist)
with pytest.raises(NotImplementedError):
dist.__rmul__(dist)
# Test addition.
allclose((dist + dist2).mean, dist.mean + dist2.mean)
allclose((dist + dist2).var, dist.var + dist2.var)
allclose((dist.__add__(b)).mean, dist.mean + b)
allclose((dist.__radd__(b)).mean, dist.mean + b)
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 predict(self, x, latent=False, return_variances=False):
"""Predict.
Args:
x (matrix): Input locations to predict at.
latent (bool, optional): Predict noiseless processes. Defaults
to `False`.
return_variances (bool, optional): Return means and variances
instead. Defaults to `False`.
Returns:
tuple[matrix]: Tuple containing means, lower 95% central credible
bound, and upper 95% central credible bound if `variances` is
`False`, and means and variances otherwise.
"""
mean, var = self.model.predict(x, latent=latent, return_variances=True)
# Pull means and variances through mixing matrix.
mean = B.dense(B.matmul(mean, self.h, tr_b=True))
var = B.dense(B.matmul(var, self.h ** 2, tr_b=True))
if not latent:
var = var + self.noise_obs
if return_variances:
return mean, var
else:
error = 1.96 * B.sqrt(var)
return mean, mean - error, mean + error
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 _check_root(a, asserted_type=object):
root = B.root(a)
# Check correctness.
approx(B.matmul(B.dense(root), B.dense(root)), B.dense(a))
# Check type.
assert isinstance(root, asserted_type)
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_concat(dense1, dense2, diag1, diag2):
with AssertDenseWarning("concatenating <dense>, <dense>, <diagonal>..."):
res = B.concat(dense1, dense2, diag1, diag2, axis=1)
dense_res = B.concat(B.dense(dense1),
B.dense(dense2),
B.dense(diag1),
B.dense(diag2),
axis=1)
approx(res, dense_res)
assert isinstance(res, Dense)
def approx(x, y, atol=0, rtol=1e-8):
"""Assert that two tensors are approximately equal.
Args:
x (tensor): First tensor.
y (tensor): Second tensor.
atol (scalar, optional): Absolute tolerance. Defaults to `0`.
rtol (scalar, optional): Relative tolerance. Defaults to `1e-8`.
"""
assert_allclose(B.dense(x), B.dense(y), atol=atol, rtol=rtol)
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 diag(a: LowRank):
if structured(a.left, a.right):
warn_upmodule(
f"Getting the diagonal of {a}: converting the factors to dense.",
category=ToDenseWarning,
)
diag_len = _diag_len(a)
left_mul = B.matmul(a.left, a.middle)
return B.sum(
B.multiply(
B.dense(left_mul)[:diag_len, :],
B.dense(a.right)[:diag_len, :]),
axis=1,
)
def predict(self, x, latent=False, return_variances=False):
"""Predict.
Args:
x (matrix): Input locations to predict at.
latent (bool, optional): Predict noiseless processes. Defaults
to `False`.
return_variances (bool, optional): Return means and variances
instead. Defaults to `False`.
Returns:
tuple: Tuple containing means, lower 95% central credible bound,
and upper 95% central credible bound if `variances` is `False`,
and means and variances otherwise.
"""
mean = B.stack(*[B.squeeze(B.dense(f.mean(x))) for f in self.fs], axis=1)
var = B.stack(*[B.squeeze(f.kernel.elwise(x)) for f in self.fs], axis=1)
if not latent:
var = var + self.noises[None, :]
if return_variances:
return mean, var
else:
error = 1.96 * B.sqrt(var)
return mean, mean - error, mean + error
def block(*rows):
"""Construct a matrix from its blocks, preserving structure when possible.
Assumes that every row has an equal number of blocks and that the sizes
of the blocks align to form a grid.
Args:
*rows (list): Rows of the block matrix.
Returns:
matrix: Assembled matrix with as much structured as possible.
"""
if len(rows) == 1 and len(rows[0]) == 1:
# There is just one block. Return it.
return rows[0][0]
res = _attempt_zero(rows)
if res is not None:
return res
res = _attempt_diagonal(rows)
if res is not None:
return res
# Could not preserve any structure. Simply concatenate them all densely.
warn_upmodule(
"Could not preserve structure in block matrix: converting to dense.",
category=ToDenseWarning,
)
return Dense(B.concat2d(*[[B.dense(x) for x in row] for row in rows]))
def diag(a, b):
# We could merge this with `block`, but `block` has a lot of overhead. It
# seems advantageous to optimise this common case.
warn_upmodule(
f"Constructing a dense block-diagonal matrix from "
f"{a} and {b}: converting to dense.",
category=ToDenseWarning,
)
a = B.dense(a)
b = B.dense(b)
dtype = B.dtype(a)
ar, ac = B.shape(a)
br, bc = B.shape(b)
return Dense(
B.concat2d([a, B.zeros(dtype, ar, bc)], [B.zeros(dtype, br, ac), b]))
def matmul(a: AbstractMatrix, b: LowerTriangular, tr_a=False, tr_b=False):
if structured(a):
warn_upmodule(
f"Matrix-multiplying {a} and {b}: converting to dense.",
category=ToDenseWarning,
)
return B.matmul(a, B.dense(b), tr_a=tr_a, tr_b=tr_b)
def power(a: AbstractMatrix, b: B.Numeric):
if structured(a):
warn_upmodule(
f"Taking an element-wise power of {a}: converting to dense.",
category=ToDenseWarning,
)
return Dense(B.power(B.dense(a), b))
def matmul(a: Diagonal, b: Kronecker, tr_a=False, tr_b=False):
warn_upmodule(
f"Cannot efficiently matrix-multiply {a} by {b}: "
f"converting the Kronecker product to dense.",
category=ToDenseWarning,
)
return B.matmul(a, B.dense(b), tr_a=tr_a, tr_b=tr_b)
def test_natural_normal():
chol = B.randn(2, 2)
dist = Normal(B.randn(2, 1), B.reg(chol @ chol.T, diag=1e-1))
nat = NaturalNormal.from_normal(dist)
# Test properties.
assert dist.dtype == nat.dtype
for name in ["dim", "mean", "var", "m2"]:
approx(getattr(dist, name), getattr(nat, name))
# Test sampling.
state = B.create_random_state(dist.dtype, seed=0)
state, sample = nat.sample(state, num=1_000_000)
emp_mean = B.mean(B.dense(sample), axis=1, squeeze=False)
emp_var = (sample - emp_mean) @ (sample - emp_mean).T / 1_000_000
approx(dist.mean, emp_mean, rtol=5e-2)
approx(dist.var, emp_var, rtol=5e-2)
# Test KL.
chol = B.randn(2, 2)
other_dist = Normal(B.randn(2, 1), B.reg(chol @ chol.T, diag=1e-2))
other_nat = NaturalNormal.from_normal(other_dist)
approx(dist.kl(other_dist), nat.kl(other_nat))
# Test log-pdf.
x = B.randn(2, 1)
approx(dist.logpdf(x), nat.logpdf(x))
def sqrt(a: AbstractMatrix):
if structured(a):
warn_upmodule(
f"Taking an element-wise square root of {a}: converting to dense.",
category=ToDenseWarning,
)
return Dense(B.sqrt(B.dense(a)))
def test_normal_logpdf(normal1):
normal1_sp = multivariate_normal(normal1.mean[:, 0], B.dense(normal1.var))
x = B.randn(3, 10)
approx(normal1.logpdf(x), normal1_sp.logpdf(x.T))
# Test the the output of `logpdf` is flattened appropriately.
assert B.shape(normal1.logpdf(B.ones(3, 1))) == ()
assert B.shape(normal1.logpdf(B.ones(3, 2))) == (2, )
def __init__(self, mat):
assert_matrix(
mat,
"Input is not a rank-2 tensor. Can only construct "
"dense matrices from rank-2 tensors.",
)
self.mat = B.dense(mat)
self.cholesky = None
def cholesky(a: Woodbury):
if a.cholesky is None:
warn_upmodule(
f"Converting {a} to dense to compute its Cholesky decomposition.",
category=ToDenseWarning,
)
a.cholesky = LowerTriangular(B.cholesky(B.reg(B.dense(a))))
return a.cholesky
def concat(*elements: AbstractMatrix, axis=0):
if structured(*elements):
elements_str = ", ".join(map(str, elements[:3]))
if len(elements) > 3:
elements_str += "..."
warn_upmodule(
f"Concatenating {elements_str}: converting to dense.",
category=ToDenseWarning,
)
return Dense(B.concat(*(B.dense(el) for el in elements), axis=axis))
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 multiply(a: LowRank, b: LowRank):
assert_compatible(a, b)
if structured(a.left, a.right, b.left, b.right):
warn_upmodule(
f"Multiplying {a} and {b}: converting factors to dense.",
category=ToDenseWarning,
)
al, am, ar = B.dense(a.left), B.dense(a.middle), B.dense(a.right)
bl, bm, br = B.dense(b.left), B.dense(b.middle), B.dense(b.right)
# Pick apart the matrices.
al, ar = B.unstack(al, axis=1), B.unstack(ar, axis=1)
bl, br = B.unstack(bl, axis=1), B.unstack(br, axis=1)
am = [B.unstack(x, axis=0) for x in B.unstack(am, axis=0)]
bm = [B.unstack(x, axis=0) for x in B.unstack(bm, axis=0)]
# Construct the factors.
left = B.stack(*[B.multiply(ali, blk) for ali in al for blk in bl], axis=1)
right = B.stack(*[B.multiply(arj, brl) for arj in ar for brl in br],
axis=1)
middle = B.stack(
*[
B.stack(*[amij * bmkl for amij in ami for bmkl in bmk], axis=0)
for ami in am for bmk in bm
],
axis=0,
)
return LowRank(left, right, middle)
def check_un_op(op, x, asserted_type=object):
"""Assert the correct of a unary operation by checking whether the
result is the same on the dense version of the argument.
Args:
op (function): Unary operation to check.
x (object): Argument.
asserted_type (type, optional): Type of result.
"""
x_dense = B.dense(x)
res = op(x)
approx(res, op(x_dense))
_assert_instance(res, asserted_type)
def pd_inv(a: Union[B.Numeric, AbstractMatrix]):
"""Invert a positive-definite matrix.
Args:
a (matrix): Positive-definite matrix to invert.
Returns:
matrix: Inverse of `a`, which is also positive definite.
"""
a = convert(a, AbstractMatrix)
# The call to `cholesky_solve` will convert the identity matrix to dense, because
# `cholesky(a)` will not have any exploitable structure. We suppress the expected
# warning by converting `B.eye(a)` to dense here already.
return B.cholesky_solve(B.cholesky(a), B.dense(B.eye(a)))
def check_bin_op(op, x, y, asserted_type=object, check_broadcasting=True):
"""Assert the correct of a binary operation by checking whether the
result is the same on the dense versions of the arguments.
Args:
op (function): Binary operation to check.
x (object): First argument.
y (object): Second argument.
asserted_type (type, optional): Type of result.
check_broadcasting (bool, optional): Check broadcasting behaviour.
"""
x_dense = B.dense(x)
y_dense = B.dense(y)
res = op(x, y)
approx(res, op(x_dense, y_dense))
with IgnoreDenseWarning():
if check_broadcasting:
approx(op(x_dense, y), op(x_dense, y_dense))
approx(op(x, y_dense), op(x_dense, y_dense))
approx(op(x_dense, y_dense), op(x_dense, y_dense))
_assert_instance(res, asserted_type)
def sample(self, state: B.RandomState, num: int = 1):
"""Sample.
Args:
state (random state): Random state.
num (int): Number of samples.
Returns:
tuple[random state, tensor]: Random state and sample.
"""
state, noise = Normal(self.prec).sample(state, num)
sample = B.cholsolve(B.chol(self.prec), B.add(noise, self.lam))
# Remove the matrix type if there is no structure. This eases working with
# JITs, which aren't happy with matrix types.
if not structured(sample):
sample = B.dense(sample)
return state, sample
def sample(self, x, latent=False):
"""Sample from the model.
Args:
x (matrix): Locations to sample at.
latent (bool, optional): Sample noiseless processes. Defaults
to `False`.
Returns:
matrix: Sample.
"""
sample = B.dense(
B.matmul(self.model.sample(x, latent=latent), self.h, tr_b=True)
)
if not latent:
sample = sample + B.sqrt(self.noise_obs) * B.randn(sample)
return sample
