setu.to_pymc#

setu.to_pymc(nle, theta, x_observed, *, conditions=None, chunk_size=None, name='nle_likelihood', dims=None)[source]#

Convert trained likelihood estimator (NLE or MixedNLE) to PyMC likelihood.

Parameters:

  • nle (LikelihoodEstimator) – Trained NLE or MixedNLE instance

  • theta (TensorVariable) – PyMC parameter variable

  • x_observed (Array) – Observed data (n_obs, x_dim) or (x_dim,)

  • conditions (Array | None) – Optional experimental conditions (n_obs, condition_dim)

  • chunk_size (int | None) – If provided, process observations in chunks of this size to reduce peak memory usage. Useful for large datasets.

  • name (str) – Name for the PyMC potential

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

Return type:

Potential

Returns:

PyMC Potential adding log p(x_observed | theta, conditions) to model

Example (without conditions):
with pm.Model():

theta = pm.Normal(“theta”, 0, 1, shape=3) likelihood = nle.to_pymc(theta, x_observed) trace = pm.sample(nuts_sampler=”blackjax”)

Example (with conditions):
with pm.Model():

theta = pm.Normal(“theta”, 0, 1, shape=2) likelihood = nle.to_pymc(theta, x_observed, conditions=trial_conditions) trace = pm.sample(nuts_sampler=”blackjax”)