From 2f5af43580759fd03ba4a1867d097bd065657a47 Mon Sep 17 00:00:00 2001 From: Abir Das Date: Sat, 25 Jul 2026 03:38:52 +0530 Subject: [PATCH] Add automatic discrete variable sampling to MCMC.get_samples() Modified MCMC.get_samples() to automatically include posterior samples of discrete latent variables that were enumerated during inference, using forward-filtering backward-sampling via infer_discrete. Previously, categorical and other discrete distributions were marginalized out during MCMC but not included in the returned samples, causing confusion for users. Changes pyro/infer/mcmc/api.py: Added include_discrete parameter to get_samples() method; capture prototype trace during sampling; use infer_discrete to sample discrete latent variables from their posterior when include_discrete=True (default) tests/infer/mcmc/test_discrete_samples.py: Added comprehensive tests for discrete variable sampling in MCMC with various configurations --- pyro/infer/mcmc/api.py | 106 ++++++++++++++++++- tests/infer/mcmc/test_discrete_samples.py | 118 ++++++++++++++++++++++ 2 files changed, 222 insertions(+), 2 deletions(-) create mode 100644 tests/infer/mcmc/test_discrete_samples.py diff --git a/pyro/infer/mcmc/api.py b/pyro/infer/mcmc/api.py index 2da7632e3c..767ea8c72b 100644 --- a/pyro/infer/mcmc/api.py +++ b/pyro/infer/mcmc/api.py @@ -474,6 +474,7 @@ def __init__( self._samples = None self._args = None self._kwargs = None + self._prototype_trace = None if save_params is not None: kernel.save_params = save_params self._validate_kernel(initial_params) @@ -554,6 +555,7 @@ def model(data): self._args, self._kwargs = args, kwargs num_samples = [0] * self.num_chains z_flat_acc = [[] for _ in range(self.num_chains)] + with optional( pyro.validation_enabled(not self.disable_validation), self.disable_validation is not None, @@ -564,6 +566,10 @@ def model(data): # requires_grad", which happens with `jit_compile` under PyTorch 1.7 args = [arg.detach() if torch.is_tensor(arg) else arg for arg in args] for x, chain_id in self.sampler.run(*args, **kwargs): + # Save prototype trace after first setup (before it gets cleaned up) + if self._prototype_trace is None and hasattr(self.kernel, '_prototype_trace'): + self._prototype_trace = self.kernel._prototype_trace + if num_samples[chain_id] == 0: num_samples[chain_id] += 1 z_structure = x @@ -606,14 +612,110 @@ def model(data): # terminate the sampler (shut down worker processes) self.sampler.terminate(True) - def get_samples(self, num_samples=None, group_by_chain=False): + def get_samples(self, num_samples=None, group_by_chain=False, include_discrete=True): """ Get samples from the MCMC run, potentially resampling with replacement. + :param int num_samples: Number of samples to return. If `None`, all samples + from the MCMC run are returned. + :param bool group_by_chain: Whether to preserve the chain dimension. If True, + all samples will have num_chains as the size of their leading dimension. + :param bool include_discrete: Whether to include samples for discrete latent + variables that were enumerated during MCMC. If True, discrete latent sites + will be sampled from their posterior using forward-filtering backward-sampling. + Defaults to True. + :return: dictionary of samples keyed by site name. + For parameter details see: :meth:`select_samples `. """ samples = self._samples - return select_samples(samples, num_samples, group_by_chain) + selected_samples = select_samples(samples, num_samples, group_by_chain) + + if not include_discrete or self._prototype_trace is None: + return selected_samples + + # Check if model has any discrete enumerable sites + has_discrete = any( + node["fn"].has_enumerate_support + for node in self._prototype_trace.nodes.values() + if node["type"] == "sample" and not node["is_observed"] + ) + + if not has_discrete: + return selected_samples + + # Sample discrete latent variables from their posterior + from pyro.infer.discrete import infer_discrete + from pyro.infer import config_enumerate + + # Determine the number of samples to generate + if num_samples is None: + if group_by_chain: + total_samples = self.num_chains * self.num_samples + else: + total_samples = self.num_chains * self.num_samples + else: + total_samples = num_samples if not group_by_chain else num_samples * self.num_chains + + # Flatten samples for processing if grouped by chain + if group_by_chain and len(selected_samples) > 0: + flattened_samples = { + k: v.reshape(-1, *v.shape[2:]) for k, v in selected_samples.items() + } + else: + flattened_samples = selected_samples + + # Generate discrete samples for each MCMC sample + discrete_samples = {} + num_flat_samples = next(iter(flattened_samples.values())).shape[0] if flattened_samples else total_samples + + for i in range(num_flat_samples): + # Get the i-th sample from continuous variables + sample_i = {k: v[i] for k, v in flattened_samples.items()} if flattened_samples else {} + + # Condition model on this sample and infer discrete variables + # Note: kernel.model might already be conditioned/wrapped, so we need the base model + conditioned_model = poutine.condition(self.kernel.model, sample_i) + configured_model = config_enumerate(conditioned_model, "parallel") + discrete_model = infer_discrete( + configured_model, + first_available_dim=-1 - (self.kernel._max_plate_nesting or 0), + temperature=1, + strict_enumeration_warning=False + ) + + # Run the model to sample discrete variables + with poutine.trace() as tr: + discrete_model(*self._args, **self._kwargs) + + # Collect discrete samples from this run + for name, node in tr.trace.nodes.items(): + if (node["type"] == "sample" and + not node["is_observed"] and + name not in flattened_samples): + if name not in discrete_samples: + discrete_samples[name] = [] + discrete_samples[name].append(node["value"]) + + # Stack discrete samples + for name in discrete_samples: + discrete_samples[name] = torch.stack(discrete_samples[name]) + + # Reshape to match group_by_chain format if needed + if group_by_chain and num_samples is None: + discrete_samples[name] = discrete_samples[name].reshape( + self.num_chains, self.num_samples, *discrete_samples[name].shape[1:] + ) + elif group_by_chain: + # For resampled case with group_by_chain, determine chain structure + samples_per_chain = num_samples + discrete_samples[name] = discrete_samples[name].reshape( + self.num_chains, samples_per_chain, *discrete_samples[name].shape[1:] + ) + + # Merge continuous and discrete samples + all_samples = {**selected_samples, **discrete_samples} + return all_samples def diagnostics(self): """ diff --git a/tests/infer/mcmc/test_discrete_samples.py b/tests/infer/mcmc/test_discrete_samples.py new file mode 100644 index 0000000000..c1164a38f6 --- /dev/null +++ b/tests/infer/mcmc/test_discrete_samples.py @@ -0,0 +1,118 @@ +# Copyright (c) 2017-2019 Uber Technologies, Inc. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +import pyro +import pyro.distributions as dist +from pyro.infer.mcmc import HMC, MCMC, NUTS +from tests.common import assert_close, assert_equal + + +@pytest.mark.parametrize("include_discrete", [True, False]) +def test_categorical_only_hmc(include_discrete): + """Test HMC with only categorical variables (issue #3368)""" + + def model(oil_pr): + oil = pyro.sample("oil", dist.Categorical(oil_pr)) + seis_dist = torch.tensor( + [[0.1, 0.3, 0.6], [0.3, 0.4, 0.3], [0.5, 0.4, 0.1]] + ) + seis = pyro.sample("seis", dist.Categorical(seis_dist[oil, :])) + return seis + + pyro.clear_param_store() + conditioned_model = pyro.condition(model, data={"seis": torch.tensor(1)}) + hmc_kernel = HMC(conditioned_model, step_size=0.9, num_steps=4) + posterior = MCMC(hmc_kernel, num_samples=10, warmup_steps=50) + posterior.run(torch.tensor([1 / 3, 1 / 3, 1 / 3])) + + samples = posterior.get_samples(include_discrete=include_discrete) + + if include_discrete: + assert "oil" in samples, "Oil samples should be present when include_discrete=True" + assert samples["oil"].shape == torch.Size([10]) + assert torch.all((samples["oil"] >= 0) & (samples["oil"] <= 2)) + else: + assert ( + len(samples) == 0 + ), "No samples expected when include_discrete=False and no continuous vars" + + +@pytest.mark.parametrize("kernel_cls", [HMC, NUTS]) +def test_mixed_continuous_discrete(kernel_cls): + """Test with both continuous and discrete variables""" + + def model(data): + p = pyro.sample("p", dist.Beta(2.0, 2.0)) + z = pyro.sample("z", dist.Categorical(torch.tensor([0.3, 0.7]))) + means = torch.tensor([0.0, 1.0]) + pyro.sample("obs", dist.Normal(means[z], 1.0), obs=data) + return z + + pyro.clear_param_store() + data = torch.tensor(0.9) + if kernel_cls == HMC: + kernel = kernel_cls(model, step_size=0.1, num_steps=10) + else: + kernel = kernel_cls(model) + posterior = MCMC(kernel, num_samples=20, warmup_steps=50) + posterior.run(data) + + samples = posterior.get_samples() + + assert "p" in samples, "Continuous variable p should be present" + assert "z" in samples, "Discrete variable z should be present" + assert samples["p"].shape == torch.Size([20]) + assert samples["z"].shape == torch.Size([20]) + assert torch.all((samples["z"] >= 0) & (samples["z"] <= 1)) + + # Test that include_discrete=False excludes discrete samples + samples_no_discrete = posterior.get_samples(include_discrete=False) + assert "p" in samples_no_discrete + assert "z" not in samples_no_discrete + + +def test_group_by_chain_with_discrete(): + """Test group_by_chain parameter with discrete variables""" + + def model(data): + z = pyro.sample("z", dist.Categorical(torch.tensor([0.5, 0.5]))) + means = torch.tensor([0.0, 1.0]) + pyro.sample("obs", dist.Normal(means[z], 1.0), obs=data) + + pyro.clear_param_store() + data = torch.tensor(0.8) + nuts_kernel = NUTS(model) + posterior = MCMC(nuts_kernel, num_samples=10, warmup_steps=20, num_chains=1) + posterior.run(data) + + samples_grouped = posterior.get_samples(group_by_chain=True) + assert "z" in samples_grouped + assert samples_grouped["z"].shape == torch.Size([1, 10]) + + samples_flat = posterior.get_samples(group_by_chain=False) + assert samples_flat["z"].shape == torch.Size([10]) + + +def test_discrete_with_plate(): + """Test discrete sampling with plate""" + + def model(data): + with pyro.plate("data", len(data)): + z = pyro.sample("z", dist.Categorical(torch.tensor([0.5, 0.5]))) + means = torch.tensor([0.0, 1.0]) + pyro.sample("obs", dist.Normal(means[z], 1.0), obs=data) + + pyro.clear_param_store() + data = torch.tensor([0.9, 1.1, 0.8]) + nuts_kernel = NUTS(model, max_plate_nesting=1) + posterior = MCMC(nuts_kernel, num_samples=15, warmup_steps=30) + posterior.run(data) + + samples = posterior.get_samples() + assert "z" in samples + # Each sample should have 3 discrete values (one per data point) + assert samples["z"].shape == torch.Size([15, 3]) + assert torch.all((samples["z"] >= 0) & (samples["z"] <= 1))