sample from distribution without storing

numpyro/infer/mcmc.py

@@ -188,11 +188,18 @@ def _sample_fn_nojit_args(state, sampler, args, kwargs):
     return (sampler.sample(state[0], args, kwargs),)
 
 
-def _collect_fn(collect_fields):
-    @cached_by(_collect_fn, collect_fields)
+def _collect_fn(collect_fields, remove_sites):
+    @cached_by(_collect_fn, collect_fields, remove_sites)
     def collect(x):
         if collect_fields:
-            return attrgetter(*collect_fields)(x[0])
+            fields = attrgetter(*collect_fields)(x[0])
+            fields = [fields] if len(collect_fields) == 1 else list(fields)
+            # fields[0] is guaranteed to be `sample_field`
+            sample_sites = fields[0].copy()
+            for site in remove_sites:
+                del sample_sites[site]
+            fields[0] = sample_sites
+            return fields[0] if len(collect_fields) == 1 else fields
         else:
             return x[0]
 
@@ -419,7 +426,7 @@ def _get_cached_init_state(self, rng_key, args, kwargs):
         except TypeError:
             return None
 
-    def _single_chain_mcmc(self, init, args, kwargs, collect_fields):
+    def _single_chain_mcmc(self, init, args, kwargs, collect_fields, remove_sites):
         rng_key, init_state, init_params = init
         # Check if _sample_fn is None, then we need to initialize the sampler.
         if init_state is None or (getattr(self.sampler, "_sample_fn", None) is None):
@@ -452,7 +459,7 @@ def _single_chain_mcmc(self, init, args, kwargs, collect_fields):
             upper_idx,
             sample_fn,
             init_val,
-            transform=_collect_fn(collect_fields),
+            transform=_collect_fn(collect_fields, remove_sites),
             progbar=self.progress_bar,
             return_last_val=True,
             thinning=self.thinning,
@@ -626,18 +633,25 @@ def run(self, rng_key, *args, extra_fields=(), init_params=None, **kwargs):
                     " as `num_chains`."
                 )
         assert isinstance(extra_fields, (tuple, list))
-        collect_fields = tuple(
-            set(
-                (self._sample_field,)
-                + tuple(self._default_fields)
-                + tuple(extra_fields)
-            )
-        )
+
+        collect_fields = {}  # we use a dictionary to ensure `sample_field` always comes first
+        remove_sites = {}
+        for field_name in (
+            (self._sample_field,) + tuple(self._default_fields) + tuple(extra_fields)
+        ):
+            if field_name.startswith(f"~{self._sample_field}."):
+                remove_sites[(field_name[len(self._sample_field) + 2 :])] = None
+            else:
+                collect_fields[field_name] = None
+        collect_fields = tuple(collect_fields.keys())
+        remove_sites = tuple(remove_sites.keys())
+
         partial_map_fn = partial(
             self._single_chain_mcmc,
             args=args,
             kwargs=kwargs,
             collect_fields=collect_fields,
+            remove_sites=remove_sites,
         )
         map_args = (rng_key, init_state, init_params)
         if self.num_chains == 1:

test/infer/test_mcmc.py

@@ -1186,3 +1186,17 @@ def model(data):
         mcmc.run(rng_key, data, extra_fields=("num_steps",))
         num_steps_list = np.array(mcmc.get_extra_fields()["num_steps"])
         assert all(step == num_steps for step in num_steps_list)
+
+
+@pytest.mark.parametrize("kernel_cls", [NUTS, BarkerMH])
+@pytest.mark.parametrize("remove_sites", [("~z.a", "~z.b"), ("~z.a", "~z.a")])
+def test_remove_sites(kernel_cls, remove_sites):
+    def model():
+        numpyro.sample("a", dist.Normal())
+        numpyro.sample("b", dist.Normal())
+
+    mcmc = MCMC(kernel_cls(model), num_warmup=10, num_samples=10)
+    mcmc.run(random.PRNGKey(0), extra_fields=remove_sites)
+    samps = mcmc.get_samples()
+
+    assert all([site[3:] not in samps for site in remove_sites])