def _to_list(x):
if B.rank(x) == 0:
return [x]
elif B.rank(x) == 1:
return list(x)
else:
raise ValueError(f'Could not convert "{x}" to a list.')
def convert(x):
if B.rank(x) == 0:
if isinstance(x, B.Number) and x == 0:
return Zero(B.dtype(x), 1, 1)
else:
return Constant(x, 1, 1)
elif B.rank(x) == 2:
return Dense(x)
else:
raise RuntimeError(
f"Cannot convert rank {B.rank(x)} input to a matrix.")
def assert_matrix(x, message):
"""Assert that a tensor is a matrix.
Args
x (tensor): Tensor that must be a matrix.
message (str): Error message to raise in the case that `x` is not a
matrix.
"""
if B.rank(x) != 2:
raise AssertionError(message)
def assert_vector(x, message):
"""Assert that a tensor is a vector.
Args
x (tensor): Tensor that must be a vector.
message (str): Error message to raise in the case that `x` is not a
vector.
"""
if B.rank(x) != 1:
raise AssertionError(message)
def wrapped_f(*args):
converted_args = ()
for i, arg in enumerate(args):
# Check if the argument is already a Pandas object.
if isinstance(arg, (pd.Series, pd.DataFrame)):
converted_args += (arg, )
continue
# It is not already a Pandas object. Attempt to convert.
if B.rank(arg) == 0:
converted_args += (arg, )
elif B.rank(arg) == 1:
converted_args += (pd.Series(arg), )
elif B.rank(arg) == 2:
converted_args += (pd.DataFrame(arg), )
else:
raise ValueError(f"Argument {i} has rank {B.rank(arg)}, which "
f"cannot be automatically converted to a "
f"Pandas object.")
return f(*converted_args)
def eig(a, compute_eigvecs=True):
if compute_eigvecs:
vals, vecs = B.eig(a, compute_eigvecs=True)
vals = B.flatten(vals)
if B.rank(vecs) == 3:
vecs = B.transpose(vecs, perm=(1, 0, 2))
vecs = B.reshape(vecs, 3, -1)
order = compute_order(vals)
return B.take(vals, order), B.abs(B.take(vecs, order, axis=1))
else:
vals = B.flatten(B.eig(a, compute_eigvecs=False))
return B.take(vals, compute_order(vals))
def test_packer():
a, b, c = B.randn(5, 10), B.randn(20), B.randn(5, 1, 15)
for packer, args in zip(
[Packer(a, b, c), Packer([a, b, c])], [(a, b, c), ((a, b, c), )]):
# Test packing.
packed = packer.pack(*args)
assert B.rank(packed) == 1
# Test unpacking.
a_, b_, c_ = packer.unpack(packed)
allclose(a, a_)
allclose(b, b_)
allclose(c, c_)
def test_ew_1d(check_lazy_shapes, batch_a, batch_b):
a = Tensor(*((3,) if batch_a else ()), 10, 1).np()
b = Tensor(*((3,) if batch_b else ()), 10, 1).np()
# Check that we can feed both rank 1 and rank 2 tensors.
for squeeze_a in [True, False]:
for squeeze_b in [True, False]:
if squeeze_a and B.rank(a) == 2:
a2 = a[..., 0]
else:
a2 = a
if squeeze_b and B.rank(b) == 2:
b2 = b[..., 0]
else:
b2 = b
approx(B.ew_dists2(a2, b2), np.abs(a - b) ** 2)
approx(B.ew_dists2(a2), np.zeros((*((3,) if batch_a else ()), 10, 1)))
approx(B.ew_dists(a2, b2), np.abs(a - b))
approx(B.ew_dists(a2), np.zeros((*((3,) if batch_a else ()), 10, 1)))
approx(B.ew_sums2(a2, b2), np.abs(a + b) ** 2)
approx(B.ew_sums2(a2), np.abs(a + a) ** 2)
approx(B.ew_sums(a2, b2), np.abs(a + b))
approx(B.ew_sums(a2), np.abs(a + a))
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 hessian(f, x):
"""Compute the Hessian of a function at a certain input.
Args:
f (function): Function to compute Hessian of.
x (column vector): Input to compute Hessian at.
differentiable (bool, optional): Make the computation of the Hessian
differentiable. Defaults to `False`.
Returns:
matrix: Hessian.
"""
if B.rank(x) != 2 or B.shape(x)[1] != 1:
raise ValueError("Input must be a column vector.")
# Use RMAD twice to preserve memory.
hess = jax.jacrev(jax.jacrev(lambda x: f(x[:, None])))(x[:, 0])
return (hess + B.transpose(hess)
) / 2 # Symmetrise to counteract numerical errors.
def unpack(package: B.Numeric, *shapes):
"""Unpack vector.
Args:
package (tensor): Tensor to unpack.
*shapes (shape): Shapes of objects to unpack.
Returns:
list[tensor]: Original objects.
"""
if B.rank(package) != 1:
raise ValueError("Package must be a vector.")
# Unpack package.
lengths = [reduce(mul, shape, 1) for shape in shapes]
i, outs = 0, []
for length, shape in zip(lengths, shapes):
outs.append(B.reshape(package[i : i + length], *shape))
i += length
return outs
def __call__(self, x):
x_rank = B.rank(x)
if x_rank == 1:
x = x[None, :, None]
elif x_rank == 2:
x = x[None, :, :]
elif x_rank == 3:
pass
else:
raise ValueError("Cannot handle inputs of rank {}.".format(x_rank))
# Apply layers one by one.
for layer in self.layers:
x = layer(x)
# Remove batch dimension, if that wasn't specified initially.
if x_rank != 3:
x = x[0, :, :]
return x
def summarise_samples(x, samples, db=False):
"""Summarise samples.
Args:
x (vector): Inputs of samples.
samples (tensor): Samples, with the first dimension corresponding
to different samples.
db (bool, optional): Convert to decibels.
Returns:
:class:`collections.namedtuple`: Named tuple containing various
statistics of the samples.
"""
x, samples = B.to_numpy(x, samples)
random_inds = np.random.permutation(B.shape(samples)[0])[:3]
def transform(x):
if db:
return 10 * np.log10(x)
else:
return x
perm = tuple(reversed(range(B.rank(samples)))) # Reverse all dimensions.
return collect(
x=B.to_numpy(x),
mean=transform(B.mean(samples, axis=0)),
var=transform(B.std(samples, axis=0))**2,
err_68_lower=transform(B.quantile(samples, 0.32, axis=0)),
err_68_upper=transform(B.quantile(samples, 1 - 0.32, axis=0)),
err_95_lower=transform(B.quantile(samples, 0.025, axis=0)),
err_95_upper=transform(B.quantile(samples, 1 - 0.025, axis=0)),
err_99_lower=transform(B.quantile(samples, 0.0015, axis=0)),
err_99_upper=transform(B.quantile(samples, 1 - 0.0015, axis=0)),
samples=transform(B.transpose(samples, perm=perm)[..., random_inds]),
all_samples=transform(B.transpose(samples, perm=perm)),
)
def __call__(self, x):
# Put the batch dimension second.
x_rank = B.rank(x)
if x_rank == 2:
x = x[:, None, :]
elif x_rank == 3:
x = B.transpose(x, perm=(1, 0, 2))
else:
raise ValueError(f"Cannot handle inputs of rank {B.rank(x)}.")
# Recurrently apply the cell.
n, batch_size, m = B.shape(x)
y0 = B.zeros(B.dtype(x), batch_size, self.cell.width)
h0 = B.tile(self.h0, batch_size, 1)
res = B.scan(self.cell, x, h0, y0)[1]
# Put the batch dimension first again.
res = B.transpose(res, perm=(1, 0, 2))
# Remove the batch dimension, if that didn't exist before.
if x_rank == 2:
res = res[0, :, :]
return res
