There is a relevant discussion here Nested inference and site names in numpyro You can play with grad(get_likelihood_from_lr)(lr) wich is numpyro-agnostic. Here get_likelihood_from_lr runs optimizer, grt logits, compute log likelihood from the logits.
1 Like