Tracing only a subset of state variables in mcmc.sample_chain

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.

Begin with the mcmc.sample_chain entry point and its trace_fn behavior described in the issue. Determine how sampling state is selected and define a way to retain alpha without storing z; done means sample_chain can return only the requested state variables while preserving existing tracing behavior.