def assert_orthogonal(x):
"""Assert that a matrix is orthogonal."""
# Check that matrix is square.
assert B.shape(x)[0] == B.shape(x)[1]
# Check that its transpose is its inverse.
approx(B.matmul(x, x, tr_a=True), B.eye(x))
def test_take_perm(dtype, check_lazy_shapes):
a = B.range(dtype, 10)
perm = B.randperm(B.dtype_int(dtype), 10)
a2 = B.take(a, perm)
assert B.dtype(perm) == B.dtype_int(dtype)
assert B.shape(a) == B.shape(a2)
assert B.dtype(a) == B.dtype(a2)
def test_multiply_kron(kron1, kron2):
if B.shape(kron1.left) == B.shape(kron2.left) and B.shape(
kron1.right) == B.shape(kron2.right):
check_bin_op(B.multiply, kron1, kron2, asserted_type=Kronecker)
else:
with pytest.raises(AssertionError):
B.matmul(kron1, kron2)
def approx(
x: Union[tuple, B.Number, B.NPNumeric],
y: Union[tuple, B.Number, B.NPNumeric],
**kw_args,
):
assert_allclose(x, y, **kw_args)
assert B.shape(x) == B.shape(y)
def test_matmul_kron(kron1, kron2):
if (B.shape(kron1.left)[1] == B.shape(kron2.left)[0]
and B.shape(kron1.right)[1] == B.shape(kron2.right)[0]):
_check_matmul(kron1, kron2, asserted_type=Kronecker)
else:
with pytest.raises(AssertionError):
B.matmul(kron1, kron2)
def assert_lower_triangular(x):
"""Assert that a matrix is lower triangular."""
# Check that matrix is square.
assert B.shape(x)[0] == B.shape(x)[1]
# Check that upper part is all zeros.
upper = x[np.triu_indices(B.shape(x)[0], k=1)]
approx(upper, B.zeros(upper))
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 multiply(a: Kronecker, b: Kronecker):
left_compatible = B.shape(a.left) == B.shape(b.left)
right_compatible = B.shape(a.right) == B.shape(b.right)
assert (
left_compatible and right_compatible
), f"Kronecker products {a} and {b} must be compatible, but they are not."
assert_compatible(a.left, b.left)
assert_compatible(a.right, b.right)
return Kronecker(B.multiply(a.left, b.left), B.multiply(a.right, b.right))
def logdet(a: Kronecker):
assert_square(
a.left, f"Left factor of {a} must be square to compute the log-determinant."
)
assert_square(
a.right,
f"Right factor of {a} must be square to compute the log-determinant.",
)
n = B.shape(a.left)[0]
m = B.shape(a.right)[0]
return m * B.logdet(a.left) + n * B.logdet(a.right)
def _attempt_zero(rows):
# Check whether the result is just zeros.
if all([all([isinstance(x, Zero) for x in row]) for row in rows]):
# Determine the resulting data type and shape.
dtype = B.dtype(rows[0][0])
grid_rows = sum([B.shape(row[0])[0] for row in rows])
grid_cols = sum([B.shape(x)[1] for x in rows[0]])
return Zero(dtype, grid_rows, grid_cols)
else:
return None
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 test_pack_unpack():
a, b, c = B.randn(5, 10), B.randn(20), B.randn(5, 1, 15)
# Test packing.
package = pack(a, b, c)
assert B.rank(package) == 1
# Test unpacking.
a2, b2, c2 = unpack(package, B.shape(a), B.shape(b), B.shape(c))
approx(a, a2)
approx(b, b2)
approx(c, c2)
# Check that the package must be a vector.
with pytest.raises(ValueError):
unpack(B.randn(2, 2), (2, 2))
def __str__(self):
rows, cols = B.shape(self)
return (
f"<upper-triangular matrix:"
f" shape={rows}x{cols},"
f" dtype={dtype_str(self)}>"
)
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_init_shape(init, shape):
if init is None and shape is None:
raise ValueError(
f"The shape must be given to automatically initialise "
f"a matrix variable.")
if shape is None:
shape = B.shape(init)
return init, shape
def f(x, option=None):
# Check that keyword arguments are passed on correctly. Only the compilation
# with PyTorch does not pass on its keyword arguments.
if not isinstance(x, B.TorchNumeric):
assert option is not None
# This requires lazy shapes:
x = B.ones(B.dtype(x), *B.shape(x))
# This requires control flow cache:
return B.cond(x[0, 0] > 0, lambda: x[1], lambda: x[2])
def matmul(a: Kronecker, b: AbstractMatrix, tr_a=False, tr_b=False):
_assert_composable(a, b, tr_a, tr_b)
a = _tr(a, tr_a)
b = _tr(b, tr_b)
# Extract the factors of the product.
left = a.left
right = a.right
l_rows, l_cols = B.shape(left)
r_rows, r_cols = B.shape(right)
cost_first_left = l_rows * l_cols * r_cols + r_rows * r_cols * l_rows
cost_first_right = r_rows * r_cols * l_cols + l_rows * l_cols * r_rows
if cost_first_left <= cost_first_right:
return _kron_id_a_b(right, _kron_a_id_b(left, b))
else:
return _kron_a_id_b(left, _kron_id_a_b(right, b))
def project(self, x, y):
"""Project data.
Args:
x (matrix): Locations of data.
y (matrix): Observations of data.
Returns:
tuple: Tuple containing the locations of the projection,
the projection, weights associated with the projection, and
a regularisation term.
"""
n = B.shape(x)[0]
available = ~B.isnan(B.to_numpy(y))
# Optimise the case where all data is available.
if B.all(available):
return self._project_pattern(x, y, (True,) * self.p)
# Extract patterns.
patterns = list(set(map(tuple, list(available))))
if len(patterns) > 30:
warnings.warn(
f"Detected {len(patterns)} patterns, which is more "
f"than 30 and can be slow.",
category=UserWarning,
)
# Per pattern, find data points that belong to it.
patterns_inds = [[] for _ in range(len(patterns))]
for i in range(n):
patterns_inds[patterns.index(tuple(available[i]))].append(i)
# Per pattern, perform the projection.
proj_xs = []
proj_ys = []
proj_ws = []
total_reg = 0
for pattern, pattern_inds in zip(patterns, patterns_inds):
proj_x, proj_y, proj_w, reg = self._project_pattern(
B.take(x, pattern_inds), B.take(y, pattern_inds), pattern
)
proj_xs.append(proj_x)
proj_ys.append(proj_y)
proj_ws.append(proj_w)
total_reg = total_reg + reg
return (
B.concat(*proj_xs, axis=0),
B.concat(*proj_ys, axis=0),
B.concat(*proj_ws, axis=0),
total_reg,
)
def constant_to_lowrank(a):
dtype = B.dtype(a)
rows, cols = B.shape(a)
middle = B.fill_diag(a.const, 1)
if rows == cols:
return LowRank(B.ones(dtype, rows, 1), middle=middle)
else:
return LowRank(B.ones(dtype, rows, 1),
B.ones(dtype, cols, 1),
middle=middle)
def __init__(
self, model, u: AbstractMatrix, s_sqrt: AbstractMatrix, noise_obs: B.Numeric
):
self.model = model
self.u = u
self.s_sqrt = s_sqrt
self.h = u @ s_sqrt
self.noise_obs = noise_obs
self.p, self.m = B.shape(u)
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 _integrate_half2(self, name1, name2, **var_map):
if self.poly.highest_power(name1) != 2 or self.poly.highest_power(
name2) != 2:
raise RuntimeError(f'Dependency on "{name1}" and {name2}" must '
f"be quadratic.")
a11 = self.poly.collect_for(Factor(name1, 2))
a22 = self.poly.collect_for(Factor(name2, 2))
a12 = self.poly.collect_for(Factor(name1,
1)).collect_for(Factor(name2, 1))
b1 = self.poly.collect_for(Factor(name1, 1)).reject(name2)
b2 = self.poly.collect_for(Factor(name2, 1)).reject(name1)
c = self.poly.reject(name1).reject(name2)
# Evaluate and scale A.
a11 = -2 * a11.eval(**var_map)
a22 = -2 * a22.eval(**var_map)
a12 = -1 * a12.eval(**var_map)
b1 = b1.eval(**var_map)
b2 = b2.eval(**var_map)
c = c.eval(**var_map)
# Determinant of A:
a_det = a11 * a22 - a12**2
# Inverse of A, which corresponds to variance of distribution after
# completing the square:
ia11 = a22 / a_det
ia12 = -a12 / a_det
ia22 = a11 / a_det
# Mean of distribution after completing the square:
mu1 = ia11 * b1 + ia12 * b2
mu2 = ia12 * b1 + ia22 * b2
# Normalise and compute CDF part.
x1 = -mu1 / safe_sqrt(ia11)
x2 = -mu2 / safe_sqrt(ia22)
rho = ia12 / safe_sqrt(ia11 * ia22)
# Evaluate CDF for all `x1` and `x2`.
orig_shape = B.shape(mu1)
num = reduce(operator.mul, orig_shape, 1)
x1 = B.reshape(x1, num)
x2 = B.reshape(x2, num)
rho = rho * B.ones(x1)
cdf_part = B.reshape(B.bvn_cdf(x1, x2, rho), *orig_shape)
# Compute exponentiated part.
quad_form = 0.5 * (ia11 * b1**2 + ia22 * b2**2 + 2 * ia12 * b1 * b2)
det_part = 2 * B.pi / safe_sqrt(a_det)
exp_part = det_part * B.exp(quad_form + c)
return self.const * cdf_part * exp_part
def _per_output(x, y, w):
p = B.shape(y)[1]
for i in range(p):
yi = y[:, i]
wi = w[:, i]
# Only return available observations.
available = ~B.isnan(yi)
yield x[available], yi[available], wi[available]
def matmul(a: LowRank, b: LowRank, tr_a=False, tr_b=False):
_assert_composable(a, b, tr_a=tr_a, tr_b=tr_b)
a = _tr(a, tr_a)
b = _tr(b, tr_b)
middle = B.matmul(a.right, b.left, tr_a=True)
middle = B.matmul(a.middle, middle, b.middle)
# Let `middle` be of type `Diagonal` if possible.
if B.shape(middle) == (1, 1):
return LowRank(a.left, b.right, Diagonal(middle[0]))
else:
return LowRank(a.left, b.right, middle)
def test_lazy_shape():
a = B.randn(2, 2)
# By default, it should be off.
assert isinstance(B.shape(a), tuple)
# Turn on.
with B.lazy_shapes():
assert isinstance(B.shape(a), Shape)
# Force lazy shapes to be off again.
B.lazy_shapes.enabled = False
assert isinstance(B.shape(a), tuple)
# Turn on again.
with B.lazy_shapes():
assert isinstance(B.shape(a), Shape)
# Should remain off.
assert isinstance(B.shape(a), tuple)
def assert_square(x, message):
"""Assert that a tensor is a square matrix.
Args
x (tensor): Tensor that must be a matrix.
message (str): Error message to raise in the case that `x` is not a
square matrix.
"""
shape = B.shape(x)
if len(shape) != 2 or shape[0] != shape[1]:
raise AssertionError(message)
def normalise_logdet(self, y):
"""Compute the log-determinant of the Jacobian of the normalisation.
Accepts multiple arguments.
Args:
y (matrix): Data that was transformed.
Returns:
scalar: Log-determinant.
"""
return -B.shape(y)[0] * B.sum(B.log(self.scale))
def sample(a, num=1): # pragma: no cover
"""Sample from covariance matrices.
Args:
a (tensor): Covariance matrix to sample from.
num (int): Number of samples.
Returns:
tensor: Samples as rank 2 column vectors.
"""
chol = B.cholesky(a)
return B.matmul(chol, B.randn(B.dtype_float(a), B.shape(chol)[1], num))
def test_recurrent():
vs = Vars(np.float32)
# Test setting the initial hidden state.
layer = Recurrent(GRU(10), B.zeros(1, 10))
layer.initialise(5, vs)
approx(layer.h0, B.zeros(1, 10))
layer = Recurrent(GRU(10))
layer.initialise(5, vs)
assert layer.h0 is not None
# Check batch consistency.
check_batch_consistency(layer, B.randn(30, 20, 5))
# Test preservation of rank upon calls.
assert B.shape(layer(B.randn(20, 5))) == (20, 10)
assert B.shape(layer(B.randn(30, 20, 5))) == (30, 20, 10)
# Check that zero-dimensional calls fail.
with pytest.raises(ValueError):
layer(0)
def test_subshape(check_lazy_shapes):
assert B.shape(B.zeros(2), 0) == 2
assert B.shape(B.zeros(2, 3, 4), 1) == 3
assert B.shape(B.zeros(2, 3, 4), 0, 2) == (2, 4)
assert B.shape(B.zeros(2, 3, 4), 0, 1, 2) == (2, 3, 4)
# Check for possible infinite recursion.
with pytest.raises(NotFoundLookupError):
B.shape(None, 1)
