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_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_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_memoize_sample(check_sample):
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'}))
x12 = j1.sample(frozenset({'x'}))
x2 = j2.sample(frozenset({'x'}))
# 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 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_eager_subs_ground(log_density):
point1 = Tensor(randn(3))
point2 = Tensor(randn(3))
d = Delta('foo', point1, log_density)
check_funsor(d(foo=point1), {}, reals(), numeric_array(float(log_density)))
check_funsor(d(foo=point2), {}, reals(), numeric_array(float('-inf')))