Save additional outputs from likelihood calls

[j-dr]

I'm currently using the NUTS sampler for an analysis that I'm conducting and there are a number of variables that are calculated in my likelihood code that I would like to save as outputs in addition to the log density and posterior samples. Is there a straight-forward way to do this with blackjax? Thanks!