setu.to_numpyro_hierarchical#
- setu.to_numpyro_hierarchical(estimator, theta, x_observed, *, conditions=None, mask=None, name='nle_likelihood_hierarchical')[source]#
Add hierarchical NLE/MixedNLE log-likelihood as a factor in NumPyro.
Handles 2-level (subject x trial) and 3-level (group x subject x trial) hierarchies using nested vmap.
- Parameters:
estimator (
LikelihoodEstimator) – Trained NLE or MixedNLE instance.theta (
Array) – Parameters from numpyro.sample(). 2-level: shape (n_subjects, theta_dim). 3-level: shape (n_groups, n_subjects, theta_dim).x_observed (
Array) – Observed data. 2-level: shape (n_subjects, n_trials, x_dim). 3-level: shape (n_groups, n_subjects, n_trials, x_dim).conditions (
Array|None) – Optional conditions matching x_observed leading dims.mask (
Array|None) – Optional boolean array matching x_observed leading dims. Where False, trial logp is zeroed (for ragged data).name (
str) – Name for the factor site.
- Return type: