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 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 test_reg(check_lazy_shapes):
old_epsilon = B.epsilon
B.epsilon = 10
a = Matrix(2, 3).np()
approx(B.reg(a, diag=None, clip=False), a + 10 * np.eye(*a.shape))
approx(B.reg(a, diag=None, clip=True), a + 10 * np.eye(*a.shape))
approx(B.reg(a, diag=1, clip=False), a + 1 * np.eye(*a.shape))
approx(B.reg(a, diag=1, clip=True), a + 10 * np.eye(*a.shape))
approx(B.reg(a, diag=100, clip=False), a + 100 * np.eye(*a.shape))
approx(B.reg(a, diag=100, clip=True), a + 100 * np.eye(*a.shape))
B.epsilon = old_epsilon
def cholesky(a: Dense):
_assert_square_cholesky(a)
if a.cholesky is None:
a.cholesky = LowerTriangular(B.cholesky(B.reg(a.mat)))
return a.cholesky
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 inverse_transform(x):
chol = B.cholesky(B.reg(x))
return B.concat(B.log(B.diag(chol)), B.tril_to_vec(chol,
offset=-1))