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')))