def test_cholesky_grad():
rng = np.random.default_rng(utt.fetch_seed())
r = rng.standard_normal((5, 5)).astype(config.floatX)
# The dots are inside the graph since Cholesky needs separable matrices
# Check the default.
utt.verify_grad(lambda r: cholesky(r.dot(r.T)), [r], 3, rng)
# Explicit lower-triangular.
utt.verify_grad(
lambda r: Cholesky(lower=True)(r.dot(r.T)),
[r],
3,
rng,
abs_tol=0.05,
rel_tol=0.05,
)
# Explicit upper-triangular.
utt.verify_grad(
lambda r: Cholesky(lower=False)(r.dot(r.T)),
[r],
3,
rng,
abs_tol=0.05,
rel_tol=0.05,
)
def test_cholesky_indef():
x = matrix()
mat = np.array([[1, 0.2], [0.2, -2]]).astype(config.floatX)
cholesky = Cholesky(lower=True, on_error="raise")
chol_f = function([x], cholesky(x))
with pytest.raises(scipy.linalg.LinAlgError):
chol_f(mat)
cholesky = Cholesky(lower=True, on_error="nan")
chol_f = function([x], cholesky(x))
assert np.all(np.isnan(chol_f(mat)))
def test_cholesky_grad_indef():
scipy = pytest.importorskip("scipy")
x = matrix()
mat = np.array([[1, 0.2], [0.2, -2]]).astype(config.floatX)
cholesky = Cholesky(lower=True, on_error="raise")
chol_f = function([x], grad(cholesky(x).sum(), [x]))
with pytest.raises(scipy.linalg.LinAlgError):
chol_f(mat)
cholesky = Cholesky(lower=True, on_error="nan")
chol_f = function([x], grad(cholesky(x).sum(), [x]))
assert np.all(np.isnan(chol_f(mat)))
def test_cholesky_grad():
pytest.importorskip("scipy")
rng = np.random.RandomState(utt.fetch_seed())
r = rng.randn(5, 5).astype(config.floatX)
# The dots are inside the graph since Cholesky needs separable matrices
# Check the default.
utt.verify_grad(lambda r: cholesky(r.dot(r.T)), [r], 3, rng)
# Explicit lower-triangular.
utt.verify_grad(lambda r: Cholesky(lower=True)(r.dot(r.T)), [r], 3, rng)
# Explicit upper-triangular.
utt.verify_grad(lambda r: Cholesky(lower=False)(r.dot(r.T)), [r], 3, rng)
def test_cholesky_and_cholesky_grad_shape():
rng = np.random.default_rng(utt.fetch_seed())
x = matrix()
for l in (cholesky(x), Cholesky(lower=True)(x), Cholesky(lower=False)(x)):
f_chol = aesara.function([x], l.shape)
g = aesara.gradient.grad(l.sum(), x)
f_cholgrad = aesara.function([x], g.shape)
topo_chol = f_chol.maker.fgraph.toposort()
topo_cholgrad = f_cholgrad.maker.fgraph.toposort()
if config.mode != "FAST_COMPILE":
assert sum([node.op.__class__ == Cholesky for node in topo_chol]) == 0
assert (
sum([node.op.__class__ == CholeskyGrad for node in topo_cholgrad]) == 0
)
for shp in [2, 3, 5]:
m = np.cov(rng.standard_normal((shp, shp + 10))).astype(config.floatX)
np.testing.assert_equal(f_chol(m), (shp, shp))
np.testing.assert_equal(f_cholgrad(m), (shp, shp))
def test_tag_solve_triangular():
cholesky_lower = Cholesky(lower=True)
cholesky_upper = Cholesky(lower=False)
A = matrix("A")
x = vector("x")
L = cholesky_lower(A)
U = cholesky_upper(A)
b1 = solve(L, x)
b2 = solve(U, x)
f = aesara.function([A, x], b1)
if config.mode != "FAST_COMPILE":
for node in f.maker.fgraph.toposort():
if isinstance(node.op, Solve):
assert node.op.assume_a != "gen" and node.op.lower
f = aesara.function([A, x], b2)
if config.mode != "FAST_COMPILE":
for node in f.maker.fgraph.toposort():
if isinstance(node.op, Solve):
assert node.op.assume_a != "gen" and not node.op.lower
def test_cholesky():
rng = np.random.default_rng(utt.fetch_seed())
r = rng.standard_normal((5, 5)).astype(config.floatX)
pd = np.dot(r, r.T)
x = matrix()
chol = cholesky(x)
# Check the default.
ch_f = function([x], chol)
check_lower_triangular(pd, ch_f)
# Explicit lower-triangular.
chol = Cholesky(lower=True)(x)
ch_f = function([x], chol)
check_lower_triangular(pd, ch_f)
# Explicit upper-triangular.
chol = Cholesky(lower=False)(x)
ch_f = function([x], chol)
check_upper_triangular(pd, ch_f)
chol = Cholesky(lower=False, on_error="nan")(x)
ch_f = function([x], chol)
check_upper_triangular(pd, ch_f)
def test_cholesky_and_cholesky_grad_shape():
pytest.importorskip("scipy")
rng = np.random.RandomState(utt.fetch_seed())
x = tensor.matrix()
for l in (cholesky(x), Cholesky(lower=True)(x), Cholesky(lower=False)(x)):
f_chol = aesara.function([x], l.shape)
g = tensor.grad(l.sum(), x)
f_cholgrad = aesara.function([x], g.shape)
topo_chol = f_chol.maker.fgraph.toposort()
topo_cholgrad = f_cholgrad.maker.fgraph.toposort()
if config.mode != "FAST_COMPILE":
assert sum([node.op.__class__ == Cholesky for node in topo_chol]) == 0
assert (
sum([node.op.__class__ == CholeskyGrad for node in topo_cholgrad]) == 0
)
for shp in [2, 3, 5]:
m = np.cov(rng.randn(shp, shp + 10)).astype(config.floatX)
np.testing.assert_equal(f_chol(m), (shp, shp))
np.testing.assert_equal(f_cholgrad(m), (shp, shp))
def test_solve_correctness(self):
scipy = pytest.importorskip("scipy")
rng = np.random.RandomState(utt.fetch_seed())
A = matrix()
b = matrix()
y = self.op(A, b)
gen_solve_func = aesara.function([A, b], y)
cholesky_lower = Cholesky(lower=True)
L = cholesky_lower(A)
y_lower = self.op(L, b)
lower_solve_func = aesara.function([L, b], y_lower)
cholesky_upper = Cholesky(lower=False)
U = cholesky_upper(A)
y_upper = self.op(U, b)
upper_solve_func = aesara.function([U, b], y_upper)
b_val = np.asarray(rng.rand(5, 1), dtype=config.floatX)
# 1-test general case
A_val = np.asarray(rng.rand(5, 5), dtype=config.floatX)
# positive definite matrix:
A_val = np.dot(A_val.transpose(), A_val)
assert np.allclose(
scipy.linalg.solve(A_val, b_val), gen_solve_func(A_val, b_val)
)
# 2-test lower traingular case
L_val = scipy.linalg.cholesky(A_val, lower=True)
assert np.allclose(
scipy.linalg.solve_triangular(L_val, b_val, lower=True),
lower_solve_func(L_val, b_val),
)
# 3-test upper traingular case
U_val = scipy.linalg.cholesky(A_val, lower=False)
assert np.allclose(
scipy.linalg.solve_triangular(U_val, b_val, lower=False),
upper_solve_func(U_val, b_val),
)
def test_cholesky():
pytest.importorskip("scipy")
rng = np.random.RandomState(utt.fetch_seed())
r = rng.randn(5, 5).astype(config.floatX)
pd = np.dot(r, r.T)
x = matrix()
chol = cholesky(x)
# Check the default.
ch_f = function([x], chol)
check_lower_triangular(pd, ch_f)
# Explicit lower-triangular.
chol = Cholesky(lower=True)(x)
ch_f = function([x], chol)
check_lower_triangular(pd, ch_f)
# Explicit upper-triangular.
chol = Cholesky(lower=False)(x)
ch_f = function([x], chol)
check_upper_triangular(pd, ch_f)
chol = Cholesky(lower=False, on_error="nan")(x)
ch_f = function([x], chol)
check_upper_triangular(pd, ch_f)
def test_magma_opt_float16(self):
ops_to_gpu = [
(MatrixInverse(), GpuMagmaMatrixInverse),
(SVD(), GpuMagmaSVD),
(QRFull(mode="reduced"), GpuMagmaQR),
(QRIncomplete(mode="r"), GpuMagmaQR),
# TODO: add support for float16 to Eigh numpy
# (Eigh(), GpuMagmaEigh),
(Cholesky(), GpuMagmaCholesky),
]
for op, gpu_op in ops_to_gpu:
A = aesara.tensor.matrix("A", dtype="float16")
fn = aesara.function([A], op(A), mode=mode_with_gpu.excluding("cusolver"))
assert any(
[isinstance(node.op, gpu_op) for node in fn.maker.fgraph.toposort()]
)
def test_correctness(self, lower):
rng = np.random.default_rng(utt.fetch_seed())
b_val = np.asarray(rng.random((5, 1)), dtype=config.floatX)
A_val = np.asarray(rng.random((5, 5)), dtype=config.floatX)
A_val = np.dot(A_val.transpose(), A_val)
C_val = scipy.linalg.cholesky(A_val, lower=lower)
A = matrix()
b = matrix()
cholesky = Cholesky(lower=lower)
C = cholesky(A)
y_lower = solve_triangular(C, b, lower=lower)
lower_solve_func = aesara.function([C, b], y_lower)
assert np.allclose(
scipy.linalg.solve_triangular(C_val, b_val, lower=lower),
lower_solve_func(C_val, b_val),
)
def MvNormalLogp():
"""Compute the log pdf of a multivariate normal distribution.
This should be used in MvNormal.logp once Theano#5908 is released.
Parameters
----------
cov: aet.matrix
The covariance matrix.
delta: aet.matrix
Array of deviations from the mean.
"""
cov = aet.matrix("cov")
cov.tag.test_value = floatX(np.eye(3))
delta = aet.matrix("delta")
delta.tag.test_value = floatX(np.zeros((2, 3)))
solve_lower = Solve(A_structure="lower_triangular")
solve_upper = Solve(A_structure="upper_triangular")
cholesky = Cholesky(lower=True, on_error="nan")
n, k = delta.shape
n, k = f(n), f(k)
chol_cov = cholesky(cov)
diag = aet.nlinalg.diag(chol_cov)
ok = aet.all(diag > 0)
chol_cov = aet.switch(ok, chol_cov, aet.fill(chol_cov, 1))
delta_trans = solve_lower(chol_cov, delta.T).T
result = n * k * aet.log(f(2) * np.pi)
result += f(2) * n * aet.sum(aet.log(diag))
result += (delta_trans**f(2)).sum()
result = f(-0.5) * result
logp = aet.switch(ok, result, -np.inf)
def dlogp(inputs, gradients):
(g_logp, ) = gradients
cov, delta = inputs
g_logp.tag.test_value = floatX(1.0)
n, k = delta.shape
chol_cov = cholesky(cov)
diag = aet.nlinalg.diag(chol_cov)
ok = aet.all(diag > 0)
chol_cov = aet.switch(ok, chol_cov, aet.fill(chol_cov, 1))
delta_trans = solve_lower(chol_cov, delta.T).T
inner = n * aet.eye(k) - aet.dot(delta_trans.T, delta_trans)
g_cov = solve_upper(chol_cov.T, inner)
g_cov = solve_upper(chol_cov.T, g_cov.T)
tau_delta = solve_upper(chol_cov.T, delta_trans.T)
g_delta = tau_delta.T
g_cov = aet.switch(ok, g_cov, -np.nan)
g_delta = aet.switch(ok, g_delta, -np.nan)
return [-0.5 * g_cov * g_logp, -g_delta * g_logp]
return OpFromGraph([cov, delta], [logp], grad_overrides=dlogp, inline=True)
