dirmeier / sbijax Goto Github PK

use arviz to compute rhat/ess/... and return inference objects instead of raw data maybe?

• add an example of SLCP with additional noise and surjective flow

Return a pytree with parameter names and the posterior samples as ni blackjax et al,

Return posterior samples with n_chain x n_samples x n_dim and not conflated

remove distrax in favor of surjectors and tfp. there is no need for this dependency

• Consistency model method (or other method that builds on diffusions)
• Flow matching posterior inference
• Quantile regression method.
• Any other notable ones?

maybe make warmup steps as argument

I think the dimensions after vmapping nuts/... with blackjax are n_draws x n_chains x n_vars which makes the ess computation incorrect (where we assumed n_chains x n_draws x n_vars)