return assume and return pytrees for everything

[dirmeier]

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