setu.nre_to_numpyro#
- setu.nre_to_numpyro(nre, theta, x_observed, *, conditions=None, name='nre_likelihood')[source]#
Add NRE log-ratio as a factor in the active NumPyro model.
- Parameters:
nre (
RatioEstimator) – Trained NRE instance.theta (
Array) – JAX array from numpyro.sample(), shape (theta_dim,).x_observed (
Array) – Observed data, shape (n_obs, x_dim) or (x_dim,).conditions (
Array|None) – Optional conditions, shape (n_obs, condition_dim).name (
str) – Name for the factor site.
- Return type: