Gathering values from the HMC kernel: How to?

Hi!

I'm new to TFP and I'm playing with HMC implementation included in TFP. Thanks for the library, it's very nice.

I stumbled upon a scenario where my kernel does some heavy computations that I would like not to redo. The computations would return a 2D matrix where the first dimension is the chains' dimension. For example:

def joint_log_prob(data, param_obs):
    heavy_computation_results = heavy_computation(data, param_obs)
    return tf.reduce_logsumexp(heavy_computation_results, axis=1)

def unnormalized_log_posterior(param_obs):
    return joint_log_prob(data, param_obs)

(...)

res = tfp.mcmc.sample_chain(
        num_results=num_steps,
        num_burnin_steps=num_burnin,
        current_state=inits,
        kernel=kernel,
        trace_fn=trace_fn
    )

Then I would like to store/access the values of heavy_computation_results for each accepted iteration.

I noticed the acceptance info is available in the previous kernel results. Then I need to either add my results to the trace_fn or store it somehow on the side. What would be the right way to achieve this goal?

Thanks!

Start by reading the sample_chain API and its trace_fn and previous kernel results references. Determine whether heavy_computation_results can be exposed for accepted iterations through existing kernel outputs or whether a new mechanism is needed. Done means the supported approach is documented or implemented with coverage for the requested per-iteration values.