def test_reduce_moment_matching_multivariate():
int_inputs = [('i', bint(4))]
real_inputs = [('x', reals(2))]
inputs = OrderedDict(int_inputs + real_inputs)
int_inputs = OrderedDict(int_inputs)
real_inputs = OrderedDict(real_inputs)
loc = numeric_array([[-10., -1.], [+10., -1.], [+10., +1.], [-10., +1.]])
precision = zeros(4, 1, 1) + ops.new_eye(loc, (2, ))
discrete = Tensor(zeros(4), int_inputs)
gaussian = Gaussian(loc, precision, inputs)
gaussian -= gaussian.log_normalizer
joint = discrete + gaussian
with interpretation(moment_matching):
actual = joint.reduce(ops.logaddexp, 'i')
assert_close(actual.reduce(ops.logaddexp), joint.reduce(ops.logaddexp))
expected_loc = zeros(2)
expected_covariance = numeric_array([[101., 0.], [0., 2.]])
expected_precision = _inverse(expected_covariance)
expected_gaussian = Gaussian(expected_loc, expected_precision, real_inputs)
expected_gaussian -= expected_gaussian.log_normalizer
expected_discrete = Tensor(ops.log(numeric_array(4.)))
expected = expected_discrete + expected_gaussian
assert_close(actual, expected, atol=1e-5, rtol=None)
def test_advanced_indexing_shape():
I, J, M, N = 4, 4, 2, 3
x = Tensor(randn((I, J)), OrderedDict([
('i', bint(I)),
('j', bint(J)),
]))
m = Tensor(numeric_array([2, 3]), OrderedDict([('m', bint(M))]), I)
n = Tensor(numeric_array([0, 1, 1]), OrderedDict([('n', bint(N))]), J)
assert x.data.shape == (I, J)
check_funsor(x(i=m), {'j': bint(J), 'm': bint(M)}, reals())
check_funsor(x(i=m, j=n), {'m': bint(M), 'n': bint(N)}, reals())
check_funsor(x(i=m, j=n, k=m), {'m': bint(M), 'n': bint(N)}, reals())
check_funsor(x(i=m, k=m), {'j': bint(J), 'm': bint(M)}, reals())
check_funsor(x(i=n), {'j': bint(J), 'n': bint(N)}, reals())
check_funsor(x(i=n, k=m), {'j': bint(J), 'n': bint(N)}, reals())
check_funsor(x(j=m), {'i': bint(I), 'm': bint(M)}, reals())
check_funsor(x(j=m, i=n), {'m': bint(M), 'n': bint(N)}, reals())
check_funsor(x(j=m, i=n, k=m), {'m': bint(M), 'n': bint(N)}, reals())
check_funsor(x(j=m, k=m), {'i': bint(I), 'm': bint(M)}, reals())
check_funsor(x(j=n), {'i': bint(I), 'n': bint(N)}, reals())
check_funsor(x(j=n, k=m), {'i': bint(I), 'n': bint(N)}, reals())
check_funsor(x(m), {'j': bint(J), 'm': bint(M)}, reals())
check_funsor(x(m, j=n), {'m': bint(M), 'n': bint(N)}, reals())
check_funsor(x(m, j=n, k=m), {'m': bint(M), 'n': bint(N)}, reals())
check_funsor(x(m, k=m), {'j': bint(J), 'm': bint(M)}, reals())
check_funsor(x(m, n), {'m': bint(M), 'n': bint(N)}, reals())
check_funsor(x(m, n, k=m), {'m': bint(M), 'n': bint(N)}, reals())
check_funsor(x(n), {'j': bint(J), 'n': bint(N)}, reals())
check_funsor(x(n, k=m), {'j': bint(J), 'n': bint(N)}, reals())
check_funsor(x(n, m), {'m': bint(M), 'n': bint(N)}, reals())
check_funsor(x(n, m, k=m), {'m': bint(M), 'n': bint(N)}, reals())
def test_advanced_indexing_shape():
I, J, M, N = 4, 4, 2, 3
x = Tensor(randn((I, J)), OrderedDict([
('i', Bint[I]),
('j', Bint[J]),
]))
m = Tensor(numeric_array([2, 3]), OrderedDict([('m', Bint[M])]), I)
n = Tensor(numeric_array([0, 1, 1]), OrderedDict([('n', Bint[N])]), J)
assert x.data.shape == (I, J)
check_funsor(x(i=m), {'j': Bint[J], 'm': Bint[M]}, Real)
check_funsor(x(i=m, j=n), {'m': Bint[M], 'n': Bint[N]}, Real)
check_funsor(x(i=m, j=n, k=m), {'m': Bint[M], 'n': Bint[N]}, Real)
check_funsor(x(i=m, k=m), {'j': Bint[J], 'm': Bint[M]}, Real)
check_funsor(x(i=n), {'j': Bint[J], 'n': Bint[N]}, Real)
check_funsor(x(i=n, k=m), {'j': Bint[J], 'n': Bint[N]}, Real)
check_funsor(x(j=m), {'i': Bint[I], 'm': Bint[M]}, Real)
check_funsor(x(j=m, i=n), {'m': Bint[M], 'n': Bint[N]}, Real)
check_funsor(x(j=m, i=n, k=m), {'m': Bint[M], 'n': Bint[N]}, Real)
check_funsor(x(j=m, k=m), {'i': Bint[I], 'm': Bint[M]}, Real)
check_funsor(x(j=n), {'i': Bint[I], 'n': Bint[N]}, Real)
check_funsor(x(j=n, k=m), {'i': Bint[I], 'n': Bint[N]}, Real)
check_funsor(x(m), {'j': Bint[J], 'm': Bint[M]}, Real)
check_funsor(x(m, j=n), {'m': Bint[M], 'n': Bint[N]}, Real)
check_funsor(x(m, j=n, k=m), {'m': Bint[M], 'n': Bint[N]}, Real)
check_funsor(x(m, k=m), {'j': Bint[J], 'm': Bint[M]}, Real)
check_funsor(x(m, n), {'m': Bint[M], 'n': Bint[N]}, Real)
check_funsor(x(m, n, k=m), {'m': Bint[M], 'n': Bint[N]}, Real)
check_funsor(x(n), {'j': Bint[J], 'n': Bint[N]}, Real)
check_funsor(x(n, k=m), {'j': Bint[J], 'n': Bint[N]}, Real)
check_funsor(x(n, m), {'m': Bint[M], 'n': Bint[N]}, Real)
check_funsor(x(n, m, k=m), {'m': Bint[M], 'n': Bint[N]}, Real)
def test_smoke(expr, expected_type):
dx = Delta('x', Tensor(randn(2, 3), OrderedDict([('i', bint(2))])))
assert isinstance(dx, Delta)
dy = Delta('y', Tensor(randn(3, 4), OrderedDict([('j', bint(3))])))
assert isinstance(dy, Delta)
t = Tensor(randn(2, 3), OrderedDict([('i', bint(2)), ('j', bint(3))]))
assert isinstance(t, Tensor)
g = Gaussian(info_vec=numeric_array([[0.0, 0.1, 0.2], [2.0, 3.0, 4.0]]),
precision=numeric_array([[[1.0, 0.1, 0.2], [0.1, 1.0, 0.3],
[0.2, 0.3, 1.0]],
[[1.0, 0.1, 0.2], [0.1, 1.0, 0.3],
[0.2, 0.3, 1.0]]]),
inputs=OrderedDict([('i', bint(2)), ('x', reals(3))]))
assert isinstance(g, Gaussian)
i0 = Number(1, 2)
assert isinstance(i0, Number)
x0 = Tensor(numeric_array([0.5, 0.6, 0.7]))
assert isinstance(x0, Tensor)
result = eval(expr)
assert isinstance(result, expected_type)
def test_mc_plate_gaussian():
log_measure = Gaussian(numeric_array([0.]), numeric_array([[1.]]),
(('loc', Real),)) + numeric_array(-0.9189)
integrand = Gaussian(randn((100, 1)) + 3., ones((100, 1, 1)),
(('data', Bint[100]), ('loc', Real)))
rng_key = None if get_backend() != 'jax' else np.array([0, 0], dtype=np.uint32)
res = Integrate(log_measure.sample('loc', rng_key=rng_key), integrand, 'loc')
res = res.reduce(ops.mul, 'data')
assert not ((res == float('inf')) | (res == float('-inf'))).any()
def test_smoke(expr, expected_type):
g1 = Gaussian(info_vec=numeric_array([[0.0, 0.1, 0.2], [2.0, 3.0, 4.0]]),
precision=numeric_array([[[1.0, 0.1, 0.2], [0.1, 1.0, 0.3],
[0.2, 0.3, 1.0]],
[[1.0, 0.1, 0.2], [0.1, 1.0, 0.3],
[0.2, 0.3, 1.0]]]),
inputs=OrderedDict([('i', bint(2)), ('x', reals(3))]))
assert isinstance(g1, Gaussian)
g2 = Gaussian(info_vec=numeric_array([[0.0, 0.1], [2.0, 3.0]]),
precision=numeric_array([[[1.0, 0.2], [0.2, 1.0]],
[[1.0, 0.2], [0.2, 1.0]]]),
inputs=OrderedDict([('i', bint(2)), ('y', reals(2))]))
assert isinstance(g2, Gaussian)
shift = Tensor(numeric_array([-1., 1.]), OrderedDict([('i', bint(2))]))
assert isinstance(shift, Tensor)
i0 = Number(1, 2)
assert isinstance(i0, Number)
x0 = Tensor(numeric_array([0.5, 0.6, 0.7]))
assert isinstance(x0, Tensor)
y0 = Tensor(numeric_array([[0.2, 0.3], [0.8, 0.9]]),
inputs=OrderedDict([('i', bint(2))]))
assert isinstance(y0, Tensor)
result = eval(expr)
assert isinstance(result, expected_type)
def test_reduce_moment_matching_univariate():
int_inputs = [('i', bint(2))]
real_inputs = [('x', reals())]
inputs = OrderedDict(int_inputs + real_inputs)
int_inputs = OrderedDict(int_inputs)
real_inputs = OrderedDict(real_inputs)
p = 0.8
t = 1.234
s1, s2, s3 = 2.0, 3.0, 4.0
loc = numeric_array([[-s1], [s1]])
precision = numeric_array([[[s2**-2]], [[s3**-2]]])
info_vec = (precision @ ops.unsqueeze(loc, -1)).squeeze(-1)
discrete = Tensor(ops.log(numeric_array([1 - p, p])) + t, int_inputs)
gaussian = Gaussian(info_vec, precision, inputs)
gaussian -= gaussian.log_normalizer
joint = discrete + gaussian
with interpretation(moment_matching):
actual = joint.reduce(ops.logaddexp, 'i')
assert_close(actual.reduce(ops.logaddexp), joint.reduce(ops.logaddexp))
expected_loc = numeric_array([(2 * p - 1) * s1])
expected_variance = (4 * p * (1 - p) * s1**2 + (1 - p) * s2**2 + p * s3**2)
expected_precision = numeric_array([[1 / expected_variance]])
expected_info_vec = (
expected_precision @ ops.unsqueeze(expected_loc, -1)).squeeze(-1)
expected_gaussian = Gaussian(expected_info_vec, expected_precision,
real_inputs)
expected_gaussian -= expected_gaussian.log_normalizer
expected_discrete = Tensor(numeric_array(t))
expected = expected_discrete + expected_gaussian
assert_close(actual, expected, atol=1e-5, rtol=None)
def test_normal_affine(expr):
scale = Tensor(numeric_array(0.3), OrderedDict())
x = Variable('x', Real)
y = Variable('y', Real)
expected = dist.Normal(x, scale, y)
actual = eval(expr)
assert isinstance(actual, Contraction)
assert dict(actual.inputs) == dict(expected.inputs), (actual.inputs, expected.inputs)
for ta, te in zip(actual.terms, expected.terms):
assert_close(ta.align(tuple(te.inputs)), te)
def test_memoize_sample(check_sample):
if get_backend() == "jax":
from funsor.jax.distributions import Normal
else:
from funsor.torch.distributions import Normal
rng_keys = (None, None, None) if get_backend() == "torch" \
else np.array([[0, 1], [0, 2], [0, 3]], dtype=np.uint32)
with memoize():
m, s = numeric_array(0.), numeric_array(1.)
j1 = Normal(m, s, 'x')
j2 = Normal(m, s, 'x')
x1 = j1.sample(frozenset({'x'}), rng_key=rng_keys[0])
x12 = j1.sample(frozenset({'x'}), rng_key=rng_keys[1])
x2 = j2.sample(frozenset({'x'}), rng_key=rng_keys[2])
# this assertion now passes
assert j1 is j2
# these assertions fail because sample is not memoized
if check_sample:
assert x1 is x12
assert x1 is x2
def numbers_to_tensors(*args):
"""
Convert :class:`~funsor.terms.Number` s to :class:`funsor.tensor.Tensor` s,
using any provided tensor as a prototype, if available.
"""
if any(isinstance(x, Number) for x in args):
prototype = get_default_prototype()
options = dict(dtype=prototype.dtype)
for x in args:
if isinstance(x, Tensor):
options = dict(dtype=x.data.dtype,
device=getattr(x.data, "device", None))
break
with ignore_jit_warnings():
args = tuple(
Tensor(numeric_array(x.data, **options), dtype=x.dtype
) if isinstance(x, Number) else x for x in args)
return args
def test_beta_bernoulli_conjugate(batch_shape):
batch_dims = ('i', 'j', 'k')[:len(batch_shape)]
inputs = OrderedDict((k, Bint[v]) for k, v in zip(batch_dims, batch_shape))
full_shape = batch_shape
prior = Variable("prior", Real)
concentration1 = Tensor(ops.exp(randn(full_shape)), inputs)
concentration0 = Tensor(ops.exp(randn(full_shape)), inputs)
latent = dist.Beta(concentration1, concentration0, value=prior)
conditional = dist.Bernoulli(probs=prior)
reduced = (latent + conditional).reduce(ops.logaddexp, set(["prior"]))
assert isinstance(reduced, dist.DirichletMultinomial)
concentration = stack((concentration0, concentration1), dim=-1)
assert_close(reduced.concentration, concentration)
assert_close(reduced.total_count, Tensor(numeric_array(1.)))
# we need lazy expression for Beta to draw samples from it
with interpretation(funsor.terms.lazy):
lazy_latent = dist.Beta(concentration1, concentration0, value=prior)
obs = Tensor(rand(batch_shape).round(), inputs)
_assert_conjugate_density_ok(latent, conditional, obs, lazy_latent=lazy_latent)
def test_eager_subs_ground(log_density):
point1 = Tensor(randn(3))
point2 = Tensor(randn(3))
d = Delta('foo', point1, log_density)
check_funsor(d(foo=point1), {}, Real, numeric_array(float(log_density)))
check_funsor(d(foo=point2), {}, Real, numeric_array(float('-inf')))
def test_delta_delta():
v = Variable('v', Reals[2])
point = Tensor(randn(2))
log_density = Tensor(numeric_array(0.5))
d = dist.Delta(point, log_density, v)
assert d is Delta('v', point, log_density)
