Tracing only a subset of state variables in mcmc.sample_chain

[martin-wiebusch-thg]

I would like to sample from a posterior distribution which depends on two variables: z and alpha. alpha is a scalar and z is a large vector of hidden variables. I am only interested in the marginal posterior distribution of alpha. When I construct a sample with mcmc.sample_chain it always gives me the full state of the chain for each element of the sample. For large sample sizes and large vectors z this can lead to memory problems. Is there a way to tell sample_chain to only return samples for alpha, but not for z? I have experimented with the trace_fn option, but with that it seems I'm only able to output additional variables like acceptance flags.