design: unify the Prior/Message hierarchy around a single-source density

[Jammy2211]

Design discussion only — no code until this is approved. This is the bundled design issue for the last two structural items of the priors/messages census (hub #1331 "parked" list): the single-source-of-truth density refactor (census C1, prompt bug/priors/12) and the Prior/Message hierarchy collapse (census C4, prompt bug/priors/13). Both were deliberately sequenced behind Phases 1–2 of the EP framework review; that gate opened in July — the EP statistics audit (#1332, findings F1–F10) and the formal autofit/graphical/README.md (#1334, 16 code-anchored equations) are the design inputs this issue builds on. Maintainer + contributor input wanted, especially anyone with exponential-family / EP background (@rhayes777).

Why now — the evidence base

Four confirmed bug families in 13 months trace to the same structural cause: each density is encoded in several places with independent conventions, split across two class hierarchies joined by __getattr__.

1. fix: log_prior_from_value sign-convention bug across Prior subclasses #1266 — log_prior_from_value sign convention drifted from the other density paths (4-year-old bug).
2. Priors & messages: 9 confirmed bugs — guidance wanted on 5 decisions #1331-04 — TruncatedNormalMessage generic-path pdf integrated to 2.27 while its direct paths were correct.
3. Priors & messages: 9 confirmed bugs — guidance wanted on 5 decisions #1331-05/07 — constants and validity conventions applied inconsistently between the two layers.
4. TransformedMessage.logpdf/pdf omit the transform Jacobian (base-space density at physical coordinates) #1498 (fresh, found 2026-08-18 while writing the test: property-based correctness sweep over every Prior subclass #1497 property sweep) — TransformedMessage.logpdf/pdf return the base-space density at physical coordinates (UniformPrior(0,1).message.pdf integrates to 1/(2√π)) while factor() is correct — a fourth divergent encoding of "the density of a transformed distribution".

The delegation surface, measured (main @ 7d4d931)

• GaussianPrior's class hierarchy defines 28 public methods; the NormalMessage it wraps defines 53; only 5 overlap (factor, log_normalisation, ndim, project, value_for).
• 48 message methods are reachable on the prior only via Prior.__getattr__ — invisible on the prior class, in stack traces, and to static analysis. That set includes log_prior_from_value, the function every MCMC/MLE search calls in its hot loop.

The consumer split, measured

Grepping call sites by package (files containing .method():

Consumed by autofit/graphical (EP) only	Consumed by mapper/non_linear (user-facing)
kl (5 files), merge (6), logpdf_gradient, calc_exact_update, has_exact_projection, from_mode, update_invalid, natural-parameter arithmetic (sum_/sub_natural_parameters, via operators)	value_for (16 files), cdf (5), project (3), log_prior_from_value (2), sample, zeros_like

The split the prompts hypothesised is real and clean: EP-specific operations are a coherent, separable set, and the user-facing set is small. Notably, almost nothing outside the messages package calls logpdf directly (1 file each in graphical and mapper) — relevant to the #1498 contract decision.

Proposed direction (sketch to react to, not a fait accompli)

One hierarchy, two functions per family, EP as a mixin:

Distribution                     (replaces both Prior and Message)
├── log_density(x, xp=np)        — THE density. The one place the math lives.
├── value_for(unit, xp=np)       — inverse CDF (no closed-form derivation from
│                                  log_density; each family implements it)
├── logpdf / log_prior_from_value — one-line aliases of log_density
│                                  (log_prior_from_value keeps the #1331 Option A
│                                  contract: constants dropped + log_normalisation())
├── Univariate families: Uniform, LogUniform, Gaussian, LogGaussian,
│                        TruncatedGaussian, Beta, Gamma
└── EPMixin                      (natural_parameters, log_partition, project, kl,
                                  merge, from_mode, sufficient stats, …)
    — applied to families that participate in factor graphs; the exact member
      list is the measured EP-only set above, cross-checked against
      autofit/graphical/README.md §§2–3 (projection, moment matching, cavity).

• width_modifier, id_, model-mapper integration move onto the family classes; __getattr__ delegation is deleted.
• Transforms: TransformedMessage becomes an internal composition detail of the transformed families, with one physical-density path (resolving TransformedMessage.logpdf/pdf omit the transform Jacobian (base-space density at physical coordinates) #1498 by construction). Whether its innards then migrate to a bijector library is prompt bug/priors/14 — that go/no-go deliberately hangs off this issue's outcome and is out of scope here.
• User-visible API (af.GaussianPrior(...), value_for, prior_passing, with_limits) is preserved; old names become deprecated aliases for one release cycle.

Staged migration (each stage independently shippable)

1. Introduce Distribution as a sibling layer implementing log_density/value_for for one family (Gaussian), with the existing classes delegating into it — no behaviour change, property tests (test: property-based correctness sweep over every Prior subclass #1497/PR test: property-based correctness sweep over every Prior subclass #1499) green throughout.
2. Migrate one family per PR (Gaussian → Uniform/LogUniform → LogGaussian → TruncatedGaussian → Beta/Gamma), moving each family's EP operations onto the mixin as it migrates.
3. Collapse TransformedMessage into the transformed families with a single physical-density path (TransformedMessage.logpdf/pdf omit the transform Jacobian (base-space density at physical coordinates) #1498 resolved by design here at the latest).
4. Deprecate Prior.__getattr__, the old Message entry points, and the duplicate density encodings; remove after one release cycle.

The #1497 property sweep is the safety net that makes this refactor tractable — every stage must keep all 134 property tests green, so the corrected definitions from #1345/#1348 cannot silently regress.

Decisions wanted

1. One hierarchy or two? Merge Prior+Message into Distribution + EP mixin (prompt 13's position), or keep two layers but with a single shared log_density (prompt 12 alone)? The measured consult split above supports the merge; the counterargument is churn in autofit/graphical internals that research: EP review phase 1 — statistics audit of autofit/graphical #1332 just verified correct.
2. TransformedMessage.logpdf/pdf omit the transform Jacobian (base-space density at physical coordinates) #1498 contract: should TransformedMessage.logpdf be the physical density (add log_det, matching the docs: formal Bayesian specification of autofit.graphical (EP review phase 2) #1334 docstring) or documented as base-space for EP arithmetic? The unified design makes this moot eventually, but the interim fix needs the call.
3. Scope of the EP mixin: the measured EP-only set, or the full exponential-family interface from graphical/README.md §3.3?
4. Sequencing vs the bijector migration (prompt 14): design it into stage 3 now, or keep it a separate go/no-go after this lands?

This is a months-long staged migration, not a single PR. Nothing starts until the direction is agreed here.

[Jammy2211]

Status: deferred (maintainer decision, 2026-08-18). The design decisions here are intentionally postponed — nothing is blocked by waiting: the census's nine bugs are fixed (#1345/#1348) and the #1499 property sweep (134 tests, merged 21288bb) locks current behaviour in as invariants, so this refactor can start whenever it earns its slot. This issue stays open as the parked design hub; the PyAutoMind tasks are moved to parked.md accordingly.

One carve-out from the deferral: decision 2 (#1498, TransformedMessage.logpdf missing the Jacobian) is a live wrong answer rather than a design preference — it can be adjudicated and fixed standalone at any time, without reopening the rest of this design.

---

Generated by Claude Code

[Jammy2211]

@matthewghgriffiths opinion on this before I vibecobe a fix appreciated!

[matthewghgriffiths]

If you're going to define a Distribution, it might be worth matching the api of distrax - https://github.com/google-deepmind/distrax/blob/main/distrax/_src/distributions/distribution.py#L49