def __init__(self, raw_dist, raw_params, expected_value_domain, xfail_reason=""):
self.raw_dist = raw_dist
self.raw_params = raw_params
self.expected_value_domain = expected_value_domain
for name, raw_param in self.raw_params:
if get_backend() != "numpy":
# we need direct access to these tensors for gradient tests
setattr(self, name, eval(raw_param))
TEST_CASES.append(self if not xfail_reason else xfail_param(self, reason=xfail_reason))
('i->', 'i'),
(',i->', 'i'),
('ai->', 'i'),
(',ai,abij->', 'ij'),
('a,ai,bij->', 'ij'),
('ai,abi,bci,cdi->', 'i'),
('aij,abij,bcij->', 'ij'),
('a,abi,bcij,cdij->', 'ij'),
]
@pytest.mark.parametrize('equation,plates', EINSUM_EXAMPLES)
@pytest.mark.parametrize('backend', ['torch', 'pyro.ops.einsum.torch_log'])
@pytest.mark.parametrize('einsum_impl,same_lazy', [
(einsum, True),
(einsum, xfail_param(False, reason="nested interpreters?")),
(naive_plated_einsum, True),
(naive_plated_einsum, False)
])
def test_einsum_complete_sharing(equation, plates, backend, einsum_impl, same_lazy):
inputs, outputs, sizes, operands, funsor_operands = make_einsum_example(equation)
with interpretation(reflect):
lazy_expr1 = einsum_impl(equation, *funsor_operands, backend=backend, plates=plates)
lazy_expr2 = lazy_expr1 if same_lazy else \
einsum_impl(equation, *funsor_operands, backend=backend, plates=plates)
with memoize():
expr1 = reinterpret(lazy_expr1)
expr2 = reinterpret(lazy_expr2)
expr3 = reinterpret(lazy_expr1)
call_count += 1
batch_shape = gate_rate.shape[:-1]
event_shape = (2, num_origins, num_destins)
gate, rate = gate_rate.reshape(batch_shape + event_shape).unbind(-3)
rate = bounded_exp(rate, bound=1e4)
gate = torch.stack((torch.zeros_like(gate), gate), dim=-1)
return gate, rate
unpack_gate_rate_0 = unpack_gate_rate[0].fn
unpack_gate_rate_1 = unpack_gate_rate[1].fn
@pytest.mark.parametrize('analytic_kl', [
False,
xfail_param(True, reason="missing pattern"),
],
ids=['monte-carlo-kl', 'analytic-kl'])
def test_bart(analytic_kl):
global call_count
call_count = 0
with interpretation(reflect):
q = Independent(
Independent(
Contraction(
ops.nullop,
ops.add,
frozenset(),
(
Tensor(
"a,b->",
"ab,a->",
"a,b,c->",
"a,a->",
"a,a,a,ab->",
"abc,bcd,cde->",
"ab,bc,cd->",
"ab,b,bc,c,cd,d->",
]
@pytest.mark.parametrize('einsum_impl', [naive_einsum, einsum])
@pytest.mark.parametrize('equation', EINSUM_EXAMPLES)
@pytest.mark.parametrize('backend', [
'pyro.ops.einsum.torch_marginal',
xfail_param('pyro.ops.einsum.torch_map', reason="wrong adjoint"),
])
def test_einsum_adjoint(einsum_impl, equation, backend):
inputs, outputs, sizes, operands, funsor_operands = make_einsum_example(
equation)
sum_op, prod_op = BACKEND_ADJOINT_OPS[backend]
with AdjointTape() as tape: # interpretation(reflect):
fwd_expr = einsum_impl(equation, *funsor_operands, backend=backend)
actuals = tape.adjoint(sum_op, prod_op, fwd_expr, funsor_operands)
for operand in operands:
pyro_require_backward(operand)
expected_out = pyro_einsum(equation,
*operands,
modulo_total=True,
def guide():
q = pyro.param("q",
torch.randn(3).exp(),
constraint=constraints.simplex)
pyro.sample("x", dist.Categorical(q))
with pyro_backend(backend):
Elbo = infer.JitTrace_ELBO if jit else infer.Trace_ELBO
elbo = Elbo(ignore_jit_warnings=True)
assert_ok(model, guide, elbo)
@pytest.mark.parametrize("backend", [
"pyro",
xfail_param("funsor", reason="missing patterns"),
])
def test_mean_field_ok(backend):
def model():
x = pyro.sample("x", dist.Normal(0., 1.))
pyro.sample("y", dist.Normal(x, 1.))
def guide():
loc = pyro.param("loc", torch.tensor(0.))
x = pyro.sample("x", dist.Normal(loc, 1.))
pyro.sample("y", dist.Normal(x, 1.))
with pyro_backend(backend):
elbo = infer.TraceMeanField_ELBO()
assert_ok(model, guide, elbo)
def backend_to_einsum_backends(backend):
backends = [
BACKEND_TO_EINSUM_BACKEND[get_backend()],
BACKEND_TO_LOGSUMEXP_BACKEND[get_backend()]
]
return backends[:1]
@pytest.mark.parametrize('equation,plates', EINSUM_EXAMPLES)
@pytest.mark.parametrize('backend', backend_to_einsum_backends(get_backend()))
@pytest.mark.parametrize(
'einsum_impl,same_lazy',
[(einsum, True),
(einsum, xfail_param(False, reason="nested interpreters?")),
(naive_plated_einsum, True), (naive_plated_einsum, False)])
def test_einsum_complete_sharing(equation, plates, backend, einsum_impl,
same_lazy):
inputs, outputs, sizes, operands, funsor_operands = make_einsum_example(
equation)
with interpretation(reflect):
lazy_expr1 = einsum_impl(equation,
*funsor_operands,
backend=backend,
plates=plates)
lazy_expr2 = lazy_expr1 if same_lazy else \
einsum_impl(equation, *funsor_operands, backend=backend, plates=plates)
with memoize():
inputs = OrderedDict((k, bint(v)) for k, v in zip(batch_dims, batch_shape))
concentration = Tensor(torch.rand(batch_shape), inputs)
rate = Tensor(torch.rand(batch_shape), inputs)
funsor_dist = (dist.Gamma if reparametrized else
dist.NonreparameterizedGamma)(concentration, rate)
_check_sample(funsor_dist,
sample_inputs,
inputs,
num_samples=200000,
atol=5e-2 if reparametrized else 1e-1)
@pytest.mark.parametrize(
"with_lazy", [True, xfail_param(False, reason="missing pattern")])
@pytest.mark.parametrize('sample_inputs', [(), ('ii', ), ('ii', 'jj'),
('ii', 'jj', 'kk')])
@pytest.mark.parametrize('batch_shape', [(), (5, ), (2, 3)], ids=str)
@pytest.mark.parametrize('reparametrized', [True, False])
def test_normal_sample(with_lazy, batch_shape, sample_inputs, reparametrized):
batch_dims = ('i', 'j', 'k')[:len(batch_shape)]
inputs = OrderedDict((k, bint(v)) for k, v in zip(batch_dims, batch_shape))
loc = Tensor(torch.randn(batch_shape), inputs)
scale = Tensor(torch.rand(batch_shape), inputs)
with interpretation(lazy if with_lazy else eager):
funsor_dist = (dist.Normal if reparametrized else
dist.NonreparameterizedNormal)(loc, scale)
_check_sample(funsor_dist,
@pytest.mark.parametrize('batch_shape', [(), (5,), (2, 3)], ids=str)
@pytest.mark.parametrize('reparametrized', [True, False])
def test_gamma_sample(batch_shape, sample_inputs, reparametrized):
batch_dims = ('i', 'j', 'k')[:len(batch_shape)]
inputs = OrderedDict((k, Bint[v]) for k, v in zip(batch_dims, batch_shape))
concentration = rand(batch_shape)
rate = rand(batch_shape)
funsor_dist_class = (dist.Gamma if reparametrized else dist.NonreparameterizedGamma)
params = (concentration, rate)
_check_sample(funsor_dist_class, params, sample_inputs, inputs, num_samples=200000,
atol=5e-2 if reparametrized else 1e-1)
@pytest.mark.parametrize("with_lazy", [True, xfail_param(False, reason="missing pattern")])
@pytest.mark.parametrize('sample_inputs', [(), ('ii',), ('ii', 'jj'), ('ii', 'jj', 'kk')])
@pytest.mark.parametrize('batch_shape', [(), (5,), (2, 3)], ids=str)
@pytest.mark.parametrize('reparametrized', [True, False])
def test_normal_sample(with_lazy, batch_shape, sample_inputs, reparametrized):
batch_dims = ('i', 'j', 'k')[:len(batch_shape)]
inputs = OrderedDict((k, Bint[v]) for k, v in zip(batch_dims, batch_shape))
loc = randn(batch_shape)
scale = rand(batch_shape)
funsor_dist_class = (dist.Normal if reparametrized else dist.NonreparameterizedNormal)
params = (loc, scale)
_check_sample(funsor_dist_class, params, sample_inputs, inputs, num_samples=200000,
atol=1e-2 if reparametrized else 1e-1, with_lazy=with_lazy)