from numpyro import optim
from numpyro.contrib.autoguide import (AutoContinuousELBO, AutoDiagonalNormal,
                                       AutoIAFNormal, AutoBNAFNormal,
                                       AutoLaplaceApproximation,
                                       AutoLowRankMultivariateNormal,
                                       AutoMultivariateNormal)
from numpyro.contrib.nn.auto_reg_nn import AutoregressiveNN
import numpyro.distributions as dist
from numpyro.distributions import constraints, transforms
from numpyro.distributions.flows import InverseAutoregressiveTransform
from numpyro.handlers import substitute
from numpyro.infer import SVI
from numpyro.infer.util import init_to_median
from numpyro.util import fori_loop

init_strategy = init_to_median(num_samples=2)


@pytest.mark.parametrize('auto_class', [
    AutoDiagonalNormal,
    AutoIAFNormal,
    AutoBNAFNormal,
    AutoMultivariateNormal,
    AutoLaplaceApproximation,
    AutoLowRankMultivariateNormal,
])
def test_beta_bernoulli(auto_class):
    data = np.array([[1.0] * 8 + [0.0] * 2, [1.0] * 4 + [0.0] * 6]).T

    def model(data):
        f = numpyro.sample('beta', dist.Beta(np.ones(2), np.ones(2)))
Example #2
0
        inv_transforms['y'].log_abs_det_jacobian(params['y'], expected_samples['y'])
    )

    base_inv_transforms = {'x': biject_to(x_prior.support), 'y': biject_to(y_prior.base_dist.support)}
    actual_samples = constrain_fn(
        handlers.seed(model, random.PRNGKey(0)), (), {}, base_inv_transforms, params)
    actual_potential_energy = potential_energy(model, (), {}, base_inv_transforms, params)

    assert_allclose(expected_samples['x'], actual_samples['x'])
    assert_allclose(expected_samples['y'], actual_samples['y'])
    assert_allclose(actual_potential_energy, expected_potential_energy)


@pytest.mark.parametrize('init_strategy', [
    init_to_feasible(),
    init_to_median(num_samples=2),
    init_to_prior(),
    init_to_uniform(),
])
def test_initialize_model_change_point(init_strategy):
    def model(data):
        alpha = 1 / np.mean(data)
        lambda1 = numpyro.sample('lambda1', dist.Exponential(alpha))
        lambda2 = numpyro.sample('lambda2', dist.Exponential(alpha))
        tau = numpyro.sample('tau', dist.Uniform(0, 1))
        lambda12 = np.where(np.arange(len(data)) < tau * len(data), lambda1, lambda2)
        numpyro.sample('obs', dist.Poisson(lambda12), obs=data)

    count_data = np.array([
        13,  24,   8,  24,   7,  35,  14,  11,  15,  11,  22,  22,  11,  57,
        11,  19,  29,   6,  19,  12,  22,  12,  18,  72,  32,   9,   7,  13,