def trace_fn(_, pkr): return (pkr.log_shrinkage_target, pkr.inner_results.log_accept_ratio, hmc_like_step_size_getter_fn(pkr))
def hmc_like_step_size_getter_fn(kernel_results): # This is here due to the circular imports. return simple_step_size_adaptation.hmc_like_step_size_getter_fn( kernel_results)