No posterior samples when using categorical distribution [bug]

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 <pyro.infer.mcmc.util.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):
         """

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