test(recon): confirm PGDReconLayerwiseLoss fresh-PGD per-site independence (#918 item 6)

[ocg-goodfire]

What

Closes item 6 of the Stage-4 adversary/source parity tracker (#918): "Confirm PGDReconLayerwiseLoss fresh-PGD per-site path by line-by-line read" — previously marked done but inferred from LOSS_PARITY_DESIGN, not read line-by-line vs torch.

Adds one pure-structure test pinning that PGDReconLayerwiseLoss factors into one fresh-PGD entry per site, each on exactly its own single site in canonical order.

Line-by-line confirmation (torch-oracle vs JAX)

Torch (PGDReconLayerwiseLoss.update → pgd_masked_recon_loss_update → _init_adv_sources / _run_pgd_loop, git tag torch-oracle):

• Loops over every target_module_paths layer; per layer calls pgd_masked_recon_loss_update with a fresh LayerRouter(layer_name=layer).
• _init_adv_sources freshly initializes sources per call → independent per layer.
• LayerRouter.get_masks = all-ones for that one layer, all-zeros elsewhere → only that layer decomposed-masked; others take the frozen path.
• _run_pgd_loop: n_steps of add_(step_size * grad.sign()) then clamp_(0,1); routing fixed across steps.
• Aggregation: Σ_layer sum_loss / Σ_layer n = mean over (layer, example).

JAX (recon.py PGDReconLayerwiseLossConfig case + train.py fresh-PGD block + adversary.init_fresh_pgd_sources):

• make_plan(per_site(site_names), AllRoutingConfig(), FreshPGDSources(...), n_samples=1) → one ReconForward per site, live_sites=(s,).
• train.py: term_key = fold_in(key, 1+term_idx); per entry routing_key, init_key = split(fold_in(term_key, entry_idx)). Distinct entry_idx per site ⇒ distinct init_key ⇒ independent fresh source.
• sign_ascend_body: clip(sources + step_size*sign(grad), 0, 1), scanned n_steps — matches torch.
• Routing drawn once per entry, shared by ascents + main forward (SPEC S24); AllRoutingConfig over the single live site == torch LayerRouter all-ones on that one layer.
• Aggregation: Σ_entry recon / n_entries (n_draws=1) = mean over sites. Equivalent to torch (equal per-entry n).

Structural note (equivalent, not a divergence): torch inits fresh sources for all modules per layer but routes only one (non-routed sources contribute nothing to the forward/gradient → no-op); JAX inits only the one live site. Loss value + ascent gradient identical; only RNG consumption differs — and torch (global RNG) vs JAX (split keys) already differ, so numeric goldens use deterministic init.

Fixture decision

No numeric torch golden exists for the fresh-PGD layerwise path (the equivalence suite covers faith/imp/stoch/ppgd only; fresh-PGD as a whole isn't in the goldens). The layerwise path shares byte-identical fresh-PGD ascent machinery with PGDReconLoss/PGDReconSubsetLoss, differing only in chunking (per_site vs one_chunk); per-entry independence is the generic plan-entry keying. Per the item's acceptance clause, I took the confirm + document branch (a torch golden would need the deprecated torch-oracle regen flow) and pinned the plan-side invariant with a pure-JAX structural test.

🤖 Generated with Claude Code