Register `Likelihood` as PyTree
If Likelihood
would be a registered PyTree, we could compile draw_sample(likelihood, primals, *a, **k)
and the data would be traced instead of inlined (I think)!
If Likelihood
would be a registered PyTree, we could compile draw_sample(likelihood, primals, *a, **k)
and the data would be traced instead of inlined (I think)!