Add factored.dag with explicit parent sets: level-wise sample (sequential cost = depth, not K)

[ealt]

Summary

independent and chain are the two extremes of a single mechanism, and everything between them is
currently unavailable. Add a factored.dag module taking explicit parent sets, whose sample
scans over topological levels and vmaps within each level. Sequential cost becomes the DAG's
depth rather than the factor count K.

Correcting my own earlier claim in this issue: I first argued the optimisation case was weak,
because I costed the compile + ghmm path — which enumerates ∏ Vᵢ Kron products regardless of
topology, so it is topology-independent. That is the flatten-to-monolithic escape hatch. The
init + composite-function path is where the factored structure is actually exploited, and there
topology matters a great deal. Thanks @ealt for the redirect.

Where topology actually enters

module	sample	update	sequential depth of sample
independent	jax.vmap	vmap	1
chain	jax.lax.scan over all K	vmap (parent token is in the observation)	K
complete	jax.lax.scan over all K	vmap + cumsum for the prefix index	K

update is vmap in all three, because by update time the whole observation is known and every
factor's control index is computable without sequencing — topology only changes how the index is
formed (roll vs cumsum). So update is not where the win is.

sample is, because factor i's variant depends on tokens sampled earlier in the same step.
independent needs no ordering and gets a full vmap; chain and complete both serialise all K.
That gap is the entire reason independent deserves its own module — and it generalises.

Proposal

Given parents: Sequence[Sequence[int]], group factors into topological levels (level 0 = roots;
level d = factors all of whose parents lie in levels < d), then:

• sample: scan over levels, vmap over the factors within a level. Every factor in a
  level has all its parents already sampled, so within-level order is irrelevant.
• update: vmap over all factors as today; the control index for factor i is the mixed-radix
  encoding of x_factors[parents[i]], gathered rather than rolled or cumsumed.
• obs_dist: same level-wise structure — chain currently scans over K factors inside a
  vmap over all V observations, so this drops from O(V·K) sequential to O(V·depth).

sigma_i has shape (∏_{p ∈ parents(i)} V_p,), which is also where the (smaller, secondary) storage
saving lives versus complete's ∏_{j<i} V_j.

What it buys

All three existing modules become special cases, which is the sign the abstraction is the right
one rather than another sibling:

topology	levels	sample cost
independent	1 level, width K	1 step
chain	K levels, width 1	K steps
fork / hub (one root, n children)	2 levels	2 steps, not n+1
balanced tree, K factors	log K levels	log K steps
complete	K levels, width 1	K steps (its prefix dependence forces a total order)

The hub case is the sharpest: sequential depth is 2 regardless of how many children, versus K
for the equivalent chain. And generate is a scan over timesteps containing a sample, so this
depth multiplies through the entire sequence-generation loop — which is the dominant cost when
building training data. It is not an academic saving.

Why the fork/hub is also worth having scientifically

Independent of performance, that topology is currently awkward to express and is the structure we most
need for belief-discovery work. In a fork A → B, A → C, siblings B and C are conditionally
independent given A but marginally dependent — the canonical conditional-independence
structure, which a chain does not contain at all (a chain has no sibling relation).

That distinction targets the one failure mode we have not solved, distinguishing a conditional link
from block membership: a chain tests whether a child is absorbed into its parent, a fork tests
whether two siblings are merged with each other. Different failures. A hub with n children also
gives n such pairs from one parent, so difficulty scales by breadth rather than chain depth.

Expressing a fork through complete works today but is wasteful and error-prone: factor 2's sigma
must be built at shape (V_0, V_1) and made constant in the V_1 axis, so you pay ∏_{j<i} V_j
control entries to encode dependence on one of them, and nothing prevents an accidental non-constant
entry silently introducing a dependence the spec never intended. It also gets chain's serial
sample rather than the 2-level one it deserves.

Suggested tests

• Equivalence: dag with parents=[[], [0], [1], …] reproduces chain exactly; parents=[[], [], …]
  reproduces independent; parents=[[], [0], [0,1], …] reproduces complete. Same samples under the
  same key, same obs_dist, same update.
• Level assignment: a cycle in parents raises; parents[i] may name any earlier factor (not only
  i-1); multiple factors may share a parent.
• The optimisation is real: assert sequential depth (number of scan iterations) equals graph
  depth, so a hub with n children does not silently regress to n+1 steps.
• Cross-check against the slow path: compile() the same spec and confirm the composite-path
  obs_dist matches the flattened GHMM's, which pins correctness against an independent implementation.