setu.DiscreteNLE

class setu.DiscreteNLE(discrete_net, theta_transform=None, condition_transform=None, x_dim=0, condition_dim=0, distribution='categorical', num_categories=None, trained=False)

Bases: Module

Discrete Neural Likelihood Estimator for categorical or count data.

Handles pure discrete observations via CategoricalMADE (categorical data) or CountMADE (count data with ZINB/NegBin/Poisson conditionals).

Use create_discrete_nle() to create and fit_discrete_nle() to train.

Parameters:

• discrete_net (CategoricalMADE | CountMADE)

• theta_transform (Standardizer | None)

• condition_transform (Standardizer | None)

• x_dim (int)

• condition_dim (int)

• distribution (str)

• num_categories (Array | None)

• trained (bool)

log_prob(x, theta, conditions=None)

Compute log p(x | theta) or log p(x | theta, conditions).

Parameters:

• x (Array) – Discrete observation (x_dim,).

• theta (Array) – Parameters (theta_dim,).

• conditions (Array | None) – Experimental conditions (condition_dim,).

Return type:

Array

Returns:

Log probability as scalar.

Raises:

ValueError – If not trained or conditions mismatch.

sample(theta, key, conditions=None, *, n_samples)

Sample from p(x | theta).

Parameters:

• theta (Array) – Single parameter vector (theta_dim,).

• key (Array) – JAX random key.

• conditions (Array | None) – Single condition vector (condition_dim,).

• n_samples (int) – Number of samples to draw.

Return type:

Array

Returns:

Samples (n_samples, x_dim).

Raises:

ValueError – If not trained or conditions mismatch.

to_pymc(theta, x_observed, *, conditions=None, chunk_size=None, name='discrete_nle_likelihood', dims=None)

Convert to PyMC likelihood.

Parameters:

• theta – PyMC parameter variable.

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

• conditions (Array | None) – Optional experimental conditions.

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

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

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

Returns:

PyMC Potential.

to_numpyro(theta, x_observed, *, conditions=None, name='discrete_nle_likelihood')

Convert to NumPyro likelihood factor.

Parameters:

• theta – JAX array from numpyro.sample().

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

• conditions (Array | None) – Optional conditions.

• name (str) – Name for the NumPyro factor site.

to_pymc_hierarchical(theta, x_observed, *, conditions=None, trial_dim=None, subject_dim=None, group_dim=None, name='discrete_nle_likelihood_hierarchical', theta_shape=None, dims=None, chunk_size=None, mask=None)

Convert to PyMC hierarchical likelihood.

Automatically handles 2-level or 3-level hierarchies via nested vmap.

Parameters:

• theta – PyMC parameter variable with dims matching hierarchy.

• x_observed – Observed data ([groups], subjects, trials, x_dim).

• conditions (Array | None) – Optional conditions matching x_observed leading dims.

• trial_dim (str | None) – Name of trial dimension (innermost level).

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

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

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

• theta_shape (tuple[int, ...] | None) – Explicit theta shape when theta.type.shape contains None (e.g. when using PyMC dims). Inferred from x_observed if not given.

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

• chunk_size (int | None) – If set, evaluate log-prob in chunks to reduce memory.

• mask (Array | None) – Boolean array matching x_observed leading dims. Where False, trial logp is zeroed (for ragged/unbalanced data).

Returns:

PyMC Potential.

to_numpyro_hierarchical(theta, x_observed, *, conditions=None, mask=None, name='discrete_nle_likelihood_hierarchical')

Convert to NumPyro hierarchical likelihood factor.

Parameters:

• theta – Parameters from numpyro.sample(). 2-level: shape (n_subjects, theta_dim). 3-level: shape (n_groups, n_subjects, theta_dim).

• x_observed – 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) – Boolean mask matching x_observed leading dims.

• name (str) – Name for the NumPyro factor site.