diff --git a/numpyro/contrib/module.py b/numpyro/contrib/module.py index f87d96d5f..a846373f7 100644 --- a/numpyro/contrib/module.py +++ b/numpyro/contrib/module.py @@ -85,6 +85,8 @@ def flax_module( # feed in dummy data to init params args = (jnp.ones(input_shape),) if input_shape is not None else args rng_key = numpyro.prng_key() + if rng_key is None: + rng_key = random.key(0) # split rng_key into a dict of rng_kind: rng_key rngs = {} if apply_rng: @@ -187,6 +189,8 @@ def haiku_module(name, nn_module, *args, input_shape=None, apply_rng=False, **kw args = (jnp.ones(input_shape),) if input_shape is not None else args # feed in dummy data to init params rng_key = numpyro.prng_key() + if rng_key is None: + rng_key = random.key(0) if with_state: nn_params, nn_state = nn_module.init(rng_key, *args, **kwargs) nn_state = dict(nn_state) diff --git a/numpyro/examples/datasets.py b/numpyro/examples/datasets.py index f04a56b3c..b86e01728 100644 --- a/numpyro/examples/datasets.py +++ b/numpyro/examples/datasets.py @@ -28,7 +28,8 @@ dset = namedtuple("dset", ["name", "urls"]) BASEBALL = dset( - "baseball", ["https://d2hg8soec8ck9v.cloudfront.net/datasets/EfronMorrisBB.txt"] + "baseball", + ["https://github.com/pyro-ppl/datasets/blob/master/EfronMorrisBB.txt?raw=true"], ) BOSTON_HOUSING = dset( @@ -37,7 +38,7 @@ ) COVTYPE = dset( - "covtype", ["https://d2hg8soec8ck9v.cloudfront.net/datasets/covtype.zip"] + "covtype", ["https://github.com/pyro-ppl/datasets/blob/master/covtype.npz?raw=true"] ) DIPPER_VOLE = dset( @@ -48,26 +49,32 @@ MNIST = dset( "mnist", [ - "https://d2hg8soec8ck9v.cloudfront.net/datasets/mnist/train-images-idx3-ubyte.gz", - "https://d2hg8soec8ck9v.cloudfront.net/datasets/mnist/train-labels-idx1-ubyte.gz", - "https://d2hg8soec8ck9v.cloudfront.net/datasets/mnist/t10k-images-idx3-ubyte.gz", - "https://d2hg8soec8ck9v.cloudfront.net/datasets/mnist/t10k-labels-idx1-ubyte.gz", + "https://github.com/pyro-ppl/datasets/blob/master/mnist/train-images-idx3-ubyte.gz?raw=true", + "https://github.com/pyro-ppl/datasets/blob/master/mnist/train-labels-idx1-ubyte.gz?raw=true", + "https://github.com/pyro-ppl/datasets/blob/master/mnist/t10k-images-idx3-ubyte.gz?raw=true", + "https://github.com/pyro-ppl/datasets/blob/master/mnist/t10k-labels-idx1-ubyte.gz?raw=true", ], ) -SP500 = dset("SP500", ["https://d2hg8soec8ck9v.cloudfront.net/datasets/SP500.csv"]) +SP500 = dset( + "SP500", ["https://github.com/pyro-ppl/datasets/blob/master/SP500.csv?raw=true"] +) UCBADMIT = dset( - "ucbadmit", ["https://d2hg8soec8ck9v.cloudfront.net/datasets/UCBadmit.csv"] + "ucbadmit", + ["https://github.com/pyro-ppl/datasets/blob/master/UCBadmit.csv?raw=true"], ) LYNXHARE = dset( - "lynxhare", ["https://d2hg8soec8ck9v.cloudfront.net/datasets/LynxHare.txt"] + "lynxhare", + ["https://github.com/pyro-ppl/datasets/blob/master/LynxHare.txt?raw=true"], ) JSB_CHORALES = dset( "jsb_chorales", - ["https://d2hg8soec8ck9v.cloudfront.net/datasets/polyphonic/jsb_chorales.pickle"], + [ + "https://github.com/pyro-ppl/datasets/blob/master/polyphonic/jsb_chorales.pickle?raw=true" + ], ) HIGGS = dset( @@ -129,7 +136,7 @@ def _load_boston_housing(): def _load_covtype(): _download(COVTYPE) - file_path = os.path.join(DATA_DIR, "covtype.zip") + file_path = os.path.join(DATA_DIR, "covtype.npz") data = np.load(file_path) return {"train": (data["data"], data["target"])} diff --git a/numpyro/infer/elbo.py b/numpyro/infer/elbo.py index f3d96d91d..b917dbcaa 100644 --- a/numpyro/infer/elbo.py +++ b/numpyro/infer/elbo.py @@ -18,9 +18,9 @@ from numpyro.handlers import replay, seed, substitute, trace from numpyro.infer.util import ( _without_rsample_stop_gradient, + compute_log_probs, get_importance_trace, is_identically_one, - log_density, ) from numpyro.ops.provenance import eval_provenance from numpyro.util import _validate_model, check_model_guide_match, find_stack_level @@ -148,12 +148,19 @@ class Trace_ELBO(ELBO): strategy, for example `jax.pmap`. :param multi_sample_guide: Whether to make an assumption that the guide proposes multiple samples. + :param sum_sites: Whether to sum the ELBO contributions from all sites or return the + contributions as a dictionary keyed by site. """ def __init__( - self, num_particles=1, vectorize_particles=True, multi_sample_guide=False + self, + num_particles: int = 1, + vectorize_particles: bool = True, + multi_sample_guide: bool = False, + sum_sites: bool = True, ): self.multi_sample_guide = multi_sample_guide + self.sum_sites = sum_sites super().__init__( num_particles=num_particles, vectorize_particles=vectorize_particles ) @@ -171,7 +178,7 @@ def single_particle_elbo(rng_key): params = param_map.copy() model_seed, guide_seed = random.split(rng_key) seeded_guide = seed(guide, guide_seed) - guide_log_density, guide_trace = log_density( + guide_log_probs, guide_trace = compute_log_probs( seeded_guide, args, kwargs, param_map ) mutable_params = { @@ -187,13 +194,13 @@ def single_particle_elbo(rng_key): if site["type"] == "plate" } - def get_model_density(key, latent): + def compute_model_log_probs(key, latent): with seed(rng_seed=key), substitute(data={**latent, **plates}): - model_log_density, model_trace = log_density( + model_log_probs, model_trace = compute_log_probs( model, args, kwargs, params ) _validate_model(model_trace, plate_warning="loose") - return model_log_density + return model_log_probs num_guide_samples = None for site in guide_trace.values(): @@ -209,15 +216,14 @@ def get_model_density(key, latent): if (site["type"] == "sample" and site["value"].size > 0) or (site["type"] == "deterministic") } - model_log_density = vmap(get_model_density)(seeds, latents) - assert model_log_density.ndim == 1 - model_log_density = model_log_density.sum(0) - # log p(z) - log q(z) - elbo_particle = (model_log_density - guide_log_density) / seeds.shape[0] + model_log_probs = vmap(compute_model_log_probs)(seeds, latents) + model_log_probs = jax.tree.map( + lambda x: jnp.sum(x, axis=0), model_log_probs + ) else: seeded_model = seed(model, model_seed) replay_model = replay(seeded_model, guide_trace) - model_log_density, model_trace = log_density( + model_log_probs, model_trace = compute_log_probs( replay_model, args, kwargs, params ) check_model_guide_match(model_trace, guide_trace) @@ -229,31 +235,43 @@ def get_model_density(key, latent): if site["type"] == "mutable" } ) - # log p(z) - log q(z) - elbo_particle = model_log_density - guide_log_density + + # log p(z) - log q(z). We cannot use jax.tree.map(jnp.subtract, ...) because + # there may be observed sites in `model_log_probs` that are not in + # `guide_log_probs` and vice versa. + union = set(model_log_probs).union(guide_log_probs) + elbo_particle = { + name: model_log_probs.get(name, 0.0) - guide_log_probs.get(name, 0.0) + for name in union + } + if self.sum_sites: + elbo_particle = sum(elbo_particle.values(), start=0.0) if mutable_params: if self.num_particles == 1: return elbo_particle, mutable_params - else: - warnings.warn( - "mutable state is currently ignored when num_particles > 1." - ) - return elbo_particle, None - else: - return elbo_particle, None + warnings.warn( + "mutable state is currently ignored when num_particles > 1." + ) + return elbo_particle, None # Return (-elbo) since by convention we do gradient descent on a loss and # the ELBO is a lower bound that needs to be maximized. if self.num_particles == 1: elbo, mutable_state = single_particle_elbo(rng_key) - return {"loss": -elbo, "mutable_state": mutable_state} + return { + "loss": jax.tree.map(jnp.negative, elbo), + "mutable_state": mutable_state, + } else: rng_keys = random.split(rng_key, self.num_particles) elbos, mutable_state = self.vectorize_particles_fn( single_particle_elbo, rng_keys ) - return {"loss": -jnp.mean(elbos), "mutable_state": mutable_state} + return { + "loss": jax.tree.map(lambda x: -jnp.mean(x), elbos), + "mutable_state": mutable_state, + } def _get_log_prob_sum(site): @@ -282,17 +300,15 @@ def _check_mean_field_requirement(model_trace, guide_trace): ] assert set(model_sites) == set(guide_sites) if model_sites != guide_sites: - ( - warnings.warn( - "Failed to verify mean field restriction on the guide. " - "To eliminate this warning, ensure model and guide sites " - "occur in the same order.\n" - + "Model sites:\n " - + "\n ".join(model_sites) - + "Guide sites:\n " - + "\n ".join(guide_sites), - stacklevel=find_stack_level(), - ), + warnings.warn( + "Failed to verify mean field restriction on the guide. " + "To eliminate this warning, ensure model and guide sites " + "occur in the same order.\n" + + "Model sites:\n " + + "\n ".join(model_sites) + + "\nGuide sites:\n " + + "\n ".join(guide_sites), + stacklevel=find_stack_level(), ) @@ -302,6 +318,15 @@ class TraceMeanField_ELBO(ELBO): ELBO estimator in NumPyro that uses analytic KL divergences when those are available. + :param num_particles: The number of particles/samples used to form the ELBO + (gradient) estimators. + :param vectorize_particles: Whether to use `jax.vmap` to compute ELBOs over the + num_particles-many particles in parallel. If False use `jax.lax.map`. + Defaults to True. You can also pass a callable to specify a custom vectorization + strategy, for example `jax.pmap`. + :param sum_sites: Whether to sum the ELBO contributions from all sites or return the + contributions as a dictionary keyed by site. + .. warning:: This estimator may give incorrect results if the mean-field condition is not satisfied. The mean field condition is a sufficient but not necessary condition for @@ -314,6 +339,15 @@ class TraceMeanField_ELBO(ELBO): dependency structures. """ + def __init__( + self, + num_particles: int = 1, + vectorize_particles: bool = True, + sum_sites: bool = True, + ) -> None: + self.sum_sites = sum_sites + super().__init__(num_particles, vectorize_particles) + def loss_with_mutable_state( self, rng_key, param_map, model, guide, *args, **kwargs ): @@ -343,50 +377,54 @@ def single_particle_elbo(rng_key): _validate_model(model_trace, plate_warning="loose") _check_mean_field_requirement(model_trace, guide_trace) - elbo_particle = 0 + elbo_particle = {} for name, model_site in model_trace.items(): if model_site["type"] == "sample": if model_site["is_observed"]: - elbo_particle = elbo_particle + _get_log_prob_sum(model_site) + elbo_particle[name] = _get_log_prob_sum(model_site) else: guide_site = guide_trace[name] try: kl_qp = kl_divergence(guide_site["fn"], model_site["fn"]) kl_qp = scale_and_mask(kl_qp, scale=guide_site["scale"]) - elbo_particle = elbo_particle - jnp.sum(kl_qp) + elbo_particle[name] = -jnp.sum(kl_qp) except NotImplementedError: - elbo_particle = ( - elbo_particle - + _get_log_prob_sum(model_site) - - _get_log_prob_sum(guide_site) - ) + elbo_particle[name] = _get_log_prob_sum( + model_site + ) - _get_log_prob_sum(guide_site) # handle auxiliary sites in the guide for name, site in guide_trace.items(): if site["type"] == "sample" and name not in model_trace: assert site["infer"].get("is_auxiliary") or site["is_observed"] - elbo_particle = elbo_particle - _get_log_prob_sum(site) + elbo_particle[name] = -_get_log_prob_sum(site) + + if self.sum_sites: + elbo_particle = sum(elbo_particle.values(), start=0.0) if mutable_params: if self.num_particles == 1: return elbo_particle, mutable_params - else: - warnings.warn( - "mutable state is currently ignored when num_particles > 1." - ) - return elbo_particle, None - else: - return elbo_particle, None + warnings.warn( + "mutable state is currently ignored when num_particles > 1." + ) + return elbo_particle, None if self.num_particles == 1: elbo, mutable_state = single_particle_elbo(rng_key) - return {"loss": -elbo, "mutable_state": mutable_state} + return { + "loss": jax.tree.map(jnp.negative, elbo), + "mutable_state": mutable_state, + } else: rng_keys = random.split(rng_key, self.num_particles) elbos, mutable_state = self.vectorize_particles_fn( single_particle_elbo, rng_keys ) - return {"loss": -jnp.mean(elbos), "mutable_state": mutable_state} + return { + "loss": jax.tree.map(lambda x: -jnp.mean(x), elbos), + "mutable_state": mutable_state, + } class RenyiELBO(ELBO): diff --git a/numpyro/util.py b/numpyro/util.py index 05e5639e8..87f34065a 100644 --- a/numpyro/util.py +++ b/numpyro/util.py @@ -438,7 +438,7 @@ def soft_vmap( # it is better to catch OOM error and reduce chunk_size by half until OOM disappears. chunk_size = batch_size if chunk_size is None else min(batch_size, chunk_size) if chunk_size > 1: - pad = chunk_size - (batch_size % chunk_size) + pad = chunk_size - batch_size % chunk_size if batch_size % chunk_size else 0 xs = jax.tree.map( lambda x: jnp.pad(x, ((0, pad),) + ((0, 0),) * (np.ndim(x) - 1)), xs ) diff --git a/test/infer/test_autoguide.py b/test/infer/test_autoguide.py index d3a02e7f4..ecf2b3c9f 100644 --- a/test/infer/test_autoguide.py +++ b/test/infer/test_autoguide.py @@ -405,7 +405,7 @@ def model(x, y): b = numpyro.sample("b", dist.Normal(0, 10).expand([3]).to_event()) mu = a + b[0] * x + b[1] * x**2 + b[2] * x**3 with numpyro.plate("N", len(x)): - numpyro.sample("y", dist.Normal(mu, 0.001), obs=y) + numpyro.sample("y", dist.Normal(mu, 0.00001), obs=y) x = random.normal(random.PRNGKey(0), (3,)) y = 1 + 2 * x + 3 * x**2 + 4 * x**3 diff --git a/test/infer/test_infer_util.py b/test/infer/test_infer_util.py index ab0133700..c9ee9e60c 100644 --- a/test/infer/test_infer_util.py +++ b/test/infer/test_infer_util.py @@ -519,3 +519,44 @@ def guide(): assert "guide-always" in called assert "model-sometimes" not in called assert "guide-sometimes" not in called + + +def test_log_likelihood_flax_nn(): + import numpy as np + + import flax.linen as nn + from jax import random + + from numpyro.contrib.module import random_flax_module + + # Simulate + rng = np.random.default_rng(99) + N = 1000 + + X = rng.normal(0, 1, size=(N, 1)) + mu = 1 + X @ np.array([0.5]) + y = rng.normal(mu, 0.5) + + # Simple linear layer + class Linear(nn.Module): + @nn.compact + def __call__(self, x): + return nn.Dense(1, use_bias=True, name="Dense")(x) + + def model(X, y=None): + sigma = numpyro.sample("sigma", dist.HalfNormal(0.1)) + priors = {"Dense.bias": dist.Normal(0, 2.5), "Dense.kernel": dist.Normal(0, 1)} + mlp = random_flax_module( + "mlp", Linear(), prior=priors, input_shape=(X.shape[1],) + ) + with numpyro.plate("data", X.shape[0]): + mu = numpyro.deterministic("mu", mlp(X).squeeze(-1)) + y = numpyro.sample("y", dist.Normal(mu, sigma), obs=y) + + # Fit model + kernel = NUTS(model, target_accept_prob=0.95) + mcmc = MCMC(kernel, num_warmup=100, num_samples=100, num_chains=1) + mcmc.run(random.PRNGKey(0), X=X, y=y) + + # run log likelihood + numpyro.infer.util.log_likelihood(model, mcmc.get_samples(), X=X, y=y) diff --git a/test/infer/test_svi.py b/test/infer/test_svi.py index 1ae8012ff..172dceb93 100644 --- a/test/infer/test_svi.py +++ b/test/infer/test_svi.py @@ -476,6 +476,46 @@ def guide(): svi.run(random.PRNGKey(0), 10) +@pytest.mark.parametrize("loss_cls", [Trace_ELBO, TraceMeanField_ELBO]) +@pytest.mark.parametrize("sum_sites", [False, True]) +@pytest.mark.parametrize("num_particles", [1, 3]) +@pytest.mark.parametrize( + "with_mutable", [False, True], ids=["with_mutable", "without_mutable"] +) +def test_elbo_by_site(loss_cls, sum_sites, num_particles, with_mutable): + if num_particles > 1 and with_mutable: + pytest.skip("Mutable state is currently ignored when num_particles > 1.") + + def model(): + x = numpyro.sample("x", dist.Normal(-1, 1)) + numpyro.sample("y", dist.Gamma(3)) + + if with_mutable: + numpyro_mutable("x1p", x + 1) + + numpyro.sample("z", dist.Normal(x, 2), obs=5) + + def guide(): + x = numpyro.sample("x", dist.Normal(2, 2)) + numpyro.sample("y", dist.LogNormal(0.1, 0.4)) + + if with_mutable: + p = numpyro_mutable("x1p", {"value": None}) + p["value"] = x + 2 + + loss = loss_cls(num_particles=num_particles, sum_sites=sum_sites) + key = random.key(9) + value = loss.loss(key, {}, model, guide) + if sum_sites: + assert value.ndim == 0 + else: + assert isinstance(value, dict) and set(value) == {"x", "y", "z"} + total = sum(value.values()) + assert_allclose( + total, loss_cls(num_particles).loss(key, {}, model, guide), rtol=1e-6 + ) + + @pytest.mark.parametrize("stable_update", [True, False]) @pytest.mark.parametrize("num_particles", [1, 10]) @pytest.mark.parametrize("elbo", [Trace_ELBO, TraceMeanField_ELBO]) diff --git a/test/test_distributions.py b/test/test_distributions.py index 03ffbf869..f499280c3 100644 --- a/test/test_distributions.py +++ b/test/test_distributions.py @@ -1586,7 +1586,7 @@ def test_entropy_categorical(): probs = _to_probs_multinom(logits) sp_dist = osp.multinomial(1, probs) for jax_dist in [dist.CategoricalLogits(logits), dist.CategoricalProbs(probs)]: - assert_allclose(jax_dist.entropy(), sp_dist.entropy()) + assert_allclose(jax_dist.entropy(), sp_dist.entropy(), rtol=1e-6, atol=1e-6) def test_mixture_log_prob(): diff --git a/test/test_example_utils.py b/test/test_example_utils.py index 65221d2d7..cf460a9d5 100644 --- a/test/test_example_utils.py +++ b/test/test_example_utils.py @@ -42,7 +42,7 @@ def mean_pixels(i, mean_pix): def test_sp500_data_load(): _, fetch = load_dataset(SP500, split="train", shuffle=False) date, value = fetch() - assert jnp.shape(date) == jnp.shape(date) == (2427,) + assert jnp.shape(date) == jnp.shape(date) == (2517,) def test_jsb_chorales(): diff --git a/test/test_transforms.py b/test/test_transforms.py index bea2c768a..74e419c6f 100644 --- a/test/test_transforms.py +++ b/test/test_transforms.py @@ -322,7 +322,7 @@ def test_bijective_transforms(transform, shape): if isinstance(transform, less_stable_transforms): atol = 1e-2 elif isinstance(transform, (L1BallTransform, RecursiveLinearTransform)): - atol = 0.1 + atol = 0.2 assert jnp.allclose(x1, x2, atol=atol) log_abs_det_jacobian = transform.log_abs_det_jacobian(x1, y)