nn(edward2): Bayesian wrappers for advanced output heads (MoG-Dense, LinearChainCRF)

## Problem / Request

After the `pyrox`/`geonnax` split, the deterministic forward kernels for the Edward2-style advanced output heads (MoG-Dense forward, LinearChainCRF forward-backward + Viterbi) live in geonnax — tracked at **jejjohnson/geonnax#12**.

This issue stays as the pyrox-side tracker for the **Bayesian wrappers** around those cores:

- **`BayesianMixtureOfGaussiansDenseHead`** — wrap `geonnax.MixtureOfGaussiansDenseHead`; register `pyrox_sample` sites on the per-component dense weights (`pi_proj`, `mu_projs[k]`, `log_sigma_projs[k]`).
- **`BayesianLinearChainCRF`** — wrap `geonnax.LinearChainCRF`; register a `pyrox_sample` site on the pairwise-potential matrix.
- **DiscreteFlow** (Bayesian) — deferred until both deterministic discrete-flow and a user request for the Bayesian variant land.

## User Story

> As a user building a multi-modal regression model on a Bayesian neural feature extractor, I want `BayesianMixtureOfGaussiansDenseHead(in=64, num_components=5, num_outputs=1)` to register Normal priors over the per-component dense weights and return a finite-mixture posterior — without composing Bayesian dense layers + manual mixture-density head logic.

> As a user fitting a Bayesian sequence tagger, I want `BayesianLinearChainCRF(num_labels=10)` to put priors on the pairwise potentials and produce calibrated label-sequence posteriors via `log_prob` + `viterbi`.

## Motivation

- **Multi-modal regression** — mixture heads on top of a BNF feature extractor (#73 demo's planned extension).
- **Sequence labelling with calibrated uncertainty** — CRF head on top of a recurrent / attention feature extractor (composes with #52).
- **Closing the gap with TFP/Edward2** for `pyrox.nn` users migrating from TF.

The bottleneck is composition with `numpyro` priors on the layer parameters — the layer math itself comes from geonnax.

## Mathematics

See **jejjohnson/geonnax#12** for the deterministic math (MoG mixture density, linear-chain forward-backward, Viterbi). On the pyrox side, the wrappers add isotropic Normal priors over the dense / pairwise weights:

$$
W \sim \mathcal{N}(0, \sigma^2),
$$

registered once per `model()` call and substituted into the geonnax core via `eqx.tree_at` before invocation — the same pattern used for `BayesianFourierNet`, `BayesianHyperLinear`, etc.

## Proposed API

```python
class BayesianMixtureOfGaussiansDenseHead(PyroxModule):
"""MoG output head with Normal priors over per-component dense weights.

Wraps ``geonnax.MixtureOfGaussiansDenseHead``. Returns mixture parameters
``(pi, mu, log_sigma)``; users wrap into a numpyro distribution
downstream (typically ``numpyro.distributions.MixtureSameFamily``).
"""
core: geonnax.MixtureOfGaussiansDenseHead
prior_scale: float = eqx.field(static=True, default=1.0)
pyrox_name: str | None = eqx.field(static=True, default=None)

@pyrox_method
def __call__(self, x: Float[Array, "*batch D_in"]) -> tuple[
Float[Array, "*batch K"],
Float[Array, "*batch K D"],
Float[Array, "*batch K D"],
]: ...

class BayesianLinearChainCRF(PyroxModule):
"""Linear-chain CRF with a Normal prior over the pairwise potentials.

Wraps ``geonnax.LinearChainCRF``. ``log_prob`` and ``viterbi`` mirror
the geonnax surface; ``__call__`` is unused (callers invoke
``log_prob`` or ``viterbi`` directly).
"""
core: geonnax.LinearChainCRF
prior_scale: float = eqx.field(static=True, default=1.0)
pyrox_name: str | None = eqx.field(static=True, default=None)

@pyrox_method
def log_prob(
self,
unary: Float[Array, "*batch T K"],
y: Int[Array, "*batch T"],
) -> Float[Array, " *batch"]: ...

@pyrox_method
def viterbi(
self,
unary: Float[Array, "*batch T K"],
) -> Int[Array, "*batch T"]: ...
```

## Example usage

```python
import jax, jax.numpy as jnp, numpyro, equinox as eqx
import geonnax
import pyrox.nn as pnn

key = jax.random.PRNGKey(0)
k_feat, k_head, k_crf = jax.random.split(key, 3)
backbone = eqx.nn.MLP(8, 64, 64, depth=3, key=k_feat)

# Bayesian MoG head on top of a feature extractor
head = pnn.BayesianMixtureOfGaussiansDenseHead(
core=geonnax.MixtureOfGaussiansDenseHead.init(64, 5, 1, key=k_head),
prior_scale=1.0,
)

def model(x, y=None):
h = jax.vmap(backbone)(x) # (N, 64) features
pi, mu, log_sigma = head(h) # Normal priors registered on dense weights
mix = numpyro.distributions.MixtureSameFamily(
numpyro.distributions.Categorical(probs=pi),
numpyro.distributions.Normal(mu[..., 0], jnp.exp(log_sigma[..., 0])),
)
numpyro.sample("y", mix, obs=y)

# Bayesian CRF head for sequence tagging
crf = pnn.BayesianLinearChainCRF(
core=geonnax.LinearChainCRF.init(num_labels=10, key=k_crf),
)

def tag_model(unary, y=None):
numpyro.factor("crf", crf.log_prob(unary, y)) # prior on pairwise potentials
```

## Tasks / sub-tasks

- [ ] Wait for **jejjohnson/geonnax#12** (deterministic cores) to land.
- [ ] `BayesianMixtureOfGaussiansDenseHead` in `src/pyrox/nn/_mixture.py`.
- [ ] `BayesianLinearChainCRF` in `src/pyrox/nn/_crf.py`.
- [ ] Re-export both from `pyrox.nn.__init__`.
- [ ] Tests:
- Both register the expected sample sites (one per dense layer for MoG; one for the pairwise matrix for CRF).
- Sites register exactly once per `model()` call (numpyro trace check).
- `BayesianMixtureOfGaussiansDenseHead` recovers a 3-mode mixture on synthetic data under SVI.
- `BayesianLinearChainCRF.log_prob` agrees with brute-force enumeration for $T \le 5$, $K \le 4$.
- `jit` / `vmap` / `grad` compatibility.
- [ ] Docs — `pyrox.nn` reference entries + "advanced Bayesian output heads" tutorial page.

## References

- Tran, D., Hoffman, M. W., Saurous, R. A. *et al.* (2018). *Simple, distributed, and accelerated probabilistic programming*. NeurIPS (Edward2 / TFP layers).
- Bishop, C. M. (1994). *Mixture density networks*. Technical Report.
- Sutton, C. & McCallum, A. (2012). *An introduction to conditional random fields*. Foundations and Trends in ML.

## Relationships

- **Blocked by**: jejjohnson/geonnax#12 (deterministic cores).
- Parent (theme epic): #46
- Related: #52