Shape conventions#

setu uses one array-shape convention everywhere: (G, S, T, *E). Getting the axis order right is what lets the same estimator serve flat, two-level, and three-level hierarchical inference without reshaping.

The (G, S, T, *E) convention#

Outermost to innermost:

Axis

Meaning

Term

G

Groups: the optional outermost hierarchy level

Subjects / Groups

S

Subjects: hierarchical level above trials

Subjects / Groups

T

Trials: the innermost repeated observations within a subject

Trials

*E

Event dimensions: the shape of a single observation (e.g. x_dim)

Groups and subjects sit outermost; trials come next; event dimensions are always last. Never put trials first.

Layouts by level#

D is theta_dim, C is condition_dim, and *E is usually a single x_dim axis.

Flat (no hierarchy): one parameter set, a batch of N observations:

theta       (D,)
x_obs       (N, *E)

2-level: S subjects, each with T trials:

theta       (S, D)
x_obs       (S, T, *E)
conditions  (S, T, C)

3-level: G groups of S subjects, each with T trials:

theta       (G, S, D)
x_obs       (G, S, T, *E)

The hierarchical export helpers (to_pymc_hierarchical, to_numpyro_hierarchical) vmap over the leading subject/group axes internally, so pass the full stacked array rather than looping in Python.

MixedNLE column order: continuous first, discrete last#

For MixedNLE, a single observation packs continuous and discrete variables into one event axis, continuous columns first, discrete columns last:

x = [ c_0, c_1, ..., c_{num_continuous-1} | d_0, d_1, ..., d_{num_discrete-1} ]
       └────── continuous (flow) ──────┘   └──── discrete (CategoricalMADE) ────┘

num_categories in tune_mixed_nle / create_mixed_nle describes the trailing discrete columns in this same order.

The argument-order footgun#

The evaluation and export entry points do not share an argument order:

Call

Argument order

nle.log_prob(x, theta)

x first, theta second

to_pymc(theta, x_observed)

theta first, x second

nle.sample(theta, key, ...)

theta first

When x_dim == theta_dim, as in the linear-Gaussian tutorial, swapping the two arguments returns a wrong but finite value instead of raising. Pass them by keyword whenever the dimensions match:

nle.log_prob(x=x_obs, theta=theta)