# 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 | {term}`Subjects / Groups` |
| `S` | Subjects: hierarchical level above trials | {term}`Subjects / Groups` |
| `T` | Trials: the innermost repeated observations within a subject | {term}`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:

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

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

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

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

```text
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 {term}`MixedNLE`, a single observation packs continuous and discrete variables
into one event axis, **continuous columns first, discrete columns last**:

```text
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:

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