setu.nre_to_pymc_hierarchical#

setu.nre_to_pymc_hierarchical(nre, theta, x_observed, *, conditions=None, theta_shape=None, chunk_size=None, mask=None, trial_dim=None, subject_dim=None, group_dim=None, name='nre_likelihood_hierarchical', dims=None)[source]#

Convert trained NRE to PyMC hierarchical likelihood.

Parameters:

  • nre (NRE) – Trained NRE instance.

  • theta (TensorVariable) – PyMC parameter variable with dims matching hierarchy.

  • x_observed (Array) – Observed data with leading dimensions for trials/subjects/groups.

  • conditions (Array | None) – Optional experimental conditions matching x_observed shape.

  • theta_shape (tuple[int, ...] | None) – Override for theta shape when theta uses dims.

  • chunk_size (int | None) – If provided, process trials in chunks.

  • mask (Array | None) – Optional boolean/float mask for ragged data.

  • trial_dim (str | None) – Name of trial dimension (innermost level). When omitted for a 2D x_observed, a single trial per subject is assumed and the trial axis is inserted automatically.

  • subject_dim (str | None) – Name of subject dimension.

  • group_dim (str | None) – Name of group dimension (optional).

  • name (str) – Name for the PyMC potential.

  • dims (tuple[str, ...] | None) – PyMC dimension names for the potential.

Return type:

Potential

Returns:

PyMC Potential adding sum of log r(x_observed, theta) to model.