def test_lowrank_attributes():
# Check default right and middle factor.
left = B.ones(3, 2)
lr = LowRank(left)
assert lr.left is left
assert lr.right is left
approx(lr.middle, B.eye(2))
assert lr.rank == 2
# Check given identical right.
lr = LowRank(left, left)
assert lr.left is left
assert lr.right is left
approx(lr.middle, B.eye(2))
assert lr.rank == 2
# Check given identical right and middle.
middle = B.ones(2, 2)
lr = LowRank(left, left, middle=middle)
assert lr.left is left
assert lr.right is left
approx(lr.middle, B.ones(2, 2))
assert lr.rank == 2
# Check given other right and middle factor.
right = 2 * B.ones(3, 2)
lr = LowRank(left, right, middle=middle)
assert lr.left is left
assert lr.right is right
approx(lr.middle, B.ones(2, 2))
assert lr.rank == 2
# Check rank in non-square case.
assert LowRank(B.ones(3, 2), B.ones(3, 2), B.ones(2, 2)).rank == 2
assert LowRank(B.ones(3, 2), B.ones(3, 1), B.ones(2, 1)).rank == 1
def test_normal_lazy_var_diag():
# If `var_diag` isn't set, the variance will be constructed to get the diagonal.
dist = Normal(lambda: B.eye(3))
approx(dist.var_diag, B.ones(3))
approx(dist._var, B.eye(3))
# If `var_diag` is set, the variance will _not_ be constructed to get the diagonal.
dist = Normal(lambda: B.eye(3), var_diag=lambda: 9)
approx(dist.var_diag, 9)
assert dist._var is None
def test_normal_lazy_nonzero_mean():
dist = Normal(lambda: B.ones(3, 1), lambda: B.eye(3))
# Nothing should be populated yet.
assert dist._mean is None
assert dist._var is None
# But they should be populated upon request.
approx(dist.mean, B.ones(3, 1))
assert dist._var is None
approx(dist.var, B.eye(3))
def test_cholesky_retry_factor(check_lazy_shapes):
# Try `cholesky_retry_factor = 1`.
B.cholesky_retry_factor = 1
B.cholesky(B.zeros(3, 3))
B.cholesky(B.zeros(3, 3) - 0.5 * B.eye(3) * B.epsilon)
with pytest.raises(np.linalg.LinAlgError):
B.cholesky(B.zeros(3, 3) - 0.5 * B.eye(3) * 10 * B.epsilon)
with pytest.raises(np.linalg.LinAlgError):
B.cholesky(B.zeros(3, 3) - 0.5 * B.eye(3) * 100 * B.epsilon)
# Try `cholesky_retry_factor = 10`.
B.cholesky_retry_factor = 10
B.cholesky(B.zeros(3, 3))
B.cholesky(B.zeros(3, 3) - 0.5 * B.eye(3) * B.epsilon)
B.cholesky(B.zeros(3, 3) - 0.5 * B.eye(3) * 10 * B.epsilon)
with pytest.raises(np.linalg.LinAlgError):
B.cholesky(B.zeros(3, 3) - 0.5 * B.eye(3) * 100 * B.epsilon)
# Try `cholesky_retry_factor = 100`.
B.cholesky_retry_factor = 100
B.cholesky(B.zeros(3, 3))
B.cholesky(B.zeros(3, 3) - 0.5 * B.eye(3) * B.epsilon)
B.cholesky(B.zeros(3, 3) - 0.5 * B.eye(3) * 10 * B.epsilon)
B.cholesky(B.zeros(3, 3) - 0.5 * B.eye(3) * 100 * B.epsilon)
# Reset the factor!
B.cholesky_retry_factor = 1
def test_normal_lazy_nonzero_mean():
dist = Normal(lambda: B.ones(3, 1), lambda: B.eye(3))
assert not dist.mean_is_zero
approx(dist._mean, B.ones(3, 1))
assert dist._var is None
approx(dist.mean, B.ones(3, 1))
assert dist._var is None
approx(dist.var, B.eye(3))
def test_normal_mean_is_zero():
# Check zero case.
dist = Normal(B.eye(3))
assert dist.mean_is_zero
approx(dist.mean, B.zeros(3, 1))
# Check another zero case.
dist = Normal(Zero(np.float32, 3, 1), B.eye(3))
assert dist.mean_is_zero
approx(dist.mean, B.zeros(3, 1))
# Check nonzero case.
assert not Normal(B.randn(3, 1), B.eye(3)).mean_is_zero
def test_normal_lazy_zero_mean():
dist = Normal(lambda: B.eye(3))
assert dist.mean_is_zero
assert dist._mean is 0
assert dist._var is None
approx(dist.mean, B.zeros(3, 1))
# At this point, the variance should be constructed, because it is used to get the
# dimensionality and data type for the mean.
assert dist._var is not None
approx(dist.var, B.eye(3))
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_normal_sampling():
for mean in [0, 1]:
dist = Normal(mean, 3 * B.eye(np.int32, 200))
# Sample without noise.
samples = dist.sample(2000)
approx(B.mean(samples), mean, atol=5e-2)
approx(B.std(samples)**2, 3, atol=5e-2)
# Sample with noise
samples = dist.sample(2000, noise=2)
approx(B.mean(samples), mean, atol=5e-2)
approx(B.std(samples)**2, 5, atol=5e-2)
def test_cholesky_solve_ut(dense_pd):
chol = B.cholesky(dense_pd)
with AssertDenseWarning(
[
"solving <upper-triangular> x = <diagonal>",
"matrix-multiplying <upper-triangular> and <lower-triangular>",
]
):
approx(
B.cholesky_solve(B.transpose(chol), B.eye(chol)),
B.inv(B.matmul(chol, chol, tr_a=True)),
)
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 test_normal_lazy_mean_var():
# The lazy `mean_var` should only be called when neither the mean nor the variance
# exists. Otherwise, it's more efficient to just construct the other one. We
# go over all branches in the `if`-statement.
dist = Normal(lambda: B.ones(3, 1), lambda: B.eye(3), mean_var=lambda: (8, 9))
approx(dist.mean_var, (8, 9))
approx(dist.mean, 8)
approx(dist.var, 9)
dist = Normal(lambda: B.ones(3, 1), lambda: B.eye(3), mean_var=lambda: (8, 9))
approx(dist.mean, B.ones(3, 1))
approx(dist.mean_var, (B.ones(3, 1), B.eye(3)))
approx(dist.var, B.eye(3))
dist = Normal(lambda: B.ones(3, 1), lambda: B.eye(3), mean_var=lambda: (8, 9))
approx(dist.var, B.eye(3))
approx(dist.mean_var, (B.ones(3, 1), B.eye(3)))
approx(dist.mean, B.ones(3, 1))
dist = Normal(lambda: B.ones(3, 1), lambda: B.eye(3), mean_var=lambda: (8, 9))
approx(dist.var, B.eye(3))
approx(dist.mean, B.ones(3, 1))
approx(dist.mean_var, (B.ones(3, 1), B.eye(3)))
def test_normal_lazy_mean_var_diag():
# The lazy `mean_var_diag` should only be called when neither the mean nor the
# diagonal of the variance exists. Otherwise, it's more efficient to just construct
# the other one. We go over all branches in the `if`-statement.
dist = Normal(lambda: B.ones(3, 1), lambda: B.eye(3), mean_var_diag=lambda: (8, 9))
approx(dist.marginals(), (8, 9))
approx(dist.mean, 8)
approx(dist.var_diag, 9)
dist = Normal(lambda: B.ones(3, 1), lambda: B.eye(3), mean_var_diag=lambda: (8, 9))
approx(dist.mean, B.ones(3, 1))
approx(dist.marginals(), (B.ones(3), B.ones(3)))
approx(dist.var_diag, B.ones(3))
dist = Normal(lambda: B.ones(3, 1), lambda: B.eye(3), mean_var_diag=lambda: (8, 9))
approx(dist.var_diag, B.ones(3))
approx(dist.marginals(), (B.ones(3), B.ones(3)))
approx(dist.mean, B.ones(3, 1))
dist = Normal(lambda: B.ones(3, 1), lambda: B.eye(3), mean_var_diag=lambda: (8, 9))
approx(dist.var_diag, B.ones(3))
approx(dist.mean, B.ones(3, 1))
approx(dist.marginals(), (B.ones(3), B.ones(3)))
def test_normal_sampling():
for mean in [0, 1]:
dist = Normal(mean, 3 * B.eye(np.int32, 200))
# Sample without noise.
samples = dist.sample(2000)
approx(B.mean(samples), mean, atol=5e-2)
approx(B.std(samples) ** 2, 3, atol=5e-2)
# Sample with noise
samples = dist.sample(2000, noise=2)
approx(B.mean(samples), mean, atol=5e-2)
approx(B.std(samples) ** 2, 5, atol=5e-2)
state, sample1 = dist.sample(B.create_random_state(B.dtype(dist), seed=0))
state, sample2 = dist.sample(B.create_random_state(B.dtype(dist), seed=0))
assert isinstance(state, B.RandomState)
approx(sample1, sample2)
def transform(x):
tril = B.vec_to_tril(x, offset=-1)
skew = tril - B.transpose(tril)
eye = B.eye(skew)
return B.solve(eye + skew, eye - skew)
def f1(x):
dists2 = (x - B.transpose(x))**2
K = B.exp(-0.5 * dists2)
K = K + B.epsilon * B.eye(t, n)
L = B.cholesky(K)
return B.matmul(L, B.ones(t, n, m))
def test_eye(dense1):
approx(B.eye(dense1), np.eye(B.shape(dense1)[0]))
assert isinstance(B.eye(dense1), Diagonal)
def test_logm(check_lazy_shapes):
mat = B.eye(3) + 0.1 * B.randn(3, 3)
check_sensitivity(logm, s_logm, (mat, ))
check_grad(logm, (mat, ))
def test_cholesky_solve_lt(dense_pd):
chol = B.cholesky(dense_pd)
with AssertDenseWarning("solving <lower-triangular> x = <diagonal>"):
approx(B.cholesky_solve(chol, B.eye(chol)), B.inv(dense_pd))
for a in Tensor(2, 2).forms():
assert "cpu" in str(B.device(a)).lower()
approx(B.to_active_device(a), a)
# Check that numbers remain unchanged.
a = 1
assert B.to_active_device(a) is a
@pytest.mark.parametrize("t", [tf.float32, torch.float32, jnp.float32])
@pytest.mark.parametrize(
"f",
[
lambda t: B.zeros(t, 2, 2),
lambda t: B.ones(t, 2, 2),
lambda t: B.eye(t, 2),
lambda t: B.linspace(t, 0, 5, 10),
lambda t: B.range(t, 10),
lambda t: B.rand(t, 10),
lambda t: B.randn(t, 10),
],
)
def test_on_device(f, t, check_lazy_shapes):
f_t = f(t) # Contruct on current and existing device.
# Set the active device to something else.
B.ActiveDevice.active_name = "previous"
# Check that explicit allocation on CPU works.
with B.on_device("cpu"):
assert B.device(f(t)) == B.device(f_t)
def generate(code):
"""Generate a random tensor of a particular type, specified with a code.
Args:
code (str): Code of the matrix.
Returns:
tensor: Random tensor.
"""
mat_code, shape_code = code.split(":")
# Parse shape.
if shape_code == "":
shape = ()
else:
shape = tuple(int(d) for d in shape_code.split(","))
if mat_code == "randn":
return B.randn(*shape)
elif mat_code == "randn_pd":
mat = B.randn(*shape)
# If it is a scalar or vector, just pointwise square it.
if len(shape) in {0, 1}:
return mat**2 + 1
else:
return B.matmul(mat, mat, tr_b=True) + B.eye(shape[0])
elif mat_code == "zero":
return Zero(B.default_dtype, *shape)
elif mat_code == "const":
return Constant(B.randn(), *shape)
elif mat_code == "const_pd":
return Constant(B.randn()**2 + 1, *shape)
elif mat_code == "lt":
mat = B.vec_to_tril(B.randn(int(0.5 * shape[0] * (shape[0] + 1))))
return LowerTriangular(mat)
elif mat_code == "lt_pd":
mat = generate(f"randn_pd:{shape[0]},{shape[0]}")
return LowerTriangular(B.cholesky(B.reg(mat)))
elif mat_code == "ut":
mat = B.vec_to_tril(B.randn(int(0.5 * shape[0] * (shape[0] + 1))))
return UpperTriangular(B.transpose(mat))
elif mat_code == "ut_pd":
mat = generate(f"randn_pd:{shape[0]},{shape[0]}")
return UpperTriangular(B.transpose(B.cholesky(B.reg(mat))))
elif mat_code == "dense":
return Dense(generate(f"randn:{shape_code}"))
elif mat_code == "dense_pd":
return Dense(generate(f"randn_pd:{shape_code}"))
elif mat_code == "diag":
return Diagonal(generate(f"randn:{shape_code}"))
elif mat_code == "diag_pd":
return Diagonal(generate(f"randn_pd:{shape_code}"))
else:
raise RuntimeError(f'Cannot parse generation code "{code}".')
def test_logm(check_lazy_shapes):
mat = B.eye(3) + 0.1 * B.randn(3, 3)
check_function(B.logm, (Tensor(mat=mat),))
def inverse_transform(x):
eye = B.eye(x)
skew = B.solve(eye + x, eye - x)
return B.tril_to_vec(skew, offset=-1)
def test_normal_dtype(normal1):
assert B.dtype(Normal(0, B.eye(3))) == np.float64
assert B.dtype(Normal(B.ones(3), B.zeros(int, 3))) == np.float64
assert B.dtype(Normal(B.ones(int, 3), B.zeros(int, 3))) == np.int64