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