gam-sae 0.3.152

//! Rung-2 two-block REML fit driver: joint activation+behavior fitting with the
//! relative block weight `λ_y` selected by REML — never a knob.
//!
//! # The objective
//!
//! The two-block model gives each output block its own Gaussian dispersion —
//! `φ_x` for the activation reconstruction, `φ_y` for the nats-unit behavior
//! tangent target — while every structural parameter (the latent coordinate
//! `t`, the gates `a`, both decoders `[B_k | C_k]`) is shared through the
//! augmented fit on `Z̃ = [Z | √λ_y·Y]` with `λ_y = φ_x/φ_y`. Profiling both
//! dispersions at a fitted state (block residuals `R_x = ‖Z − Ẑ‖²`,
//! `R_y = ‖Y − Ŷ‖²` in *unscaled* units) makes the criterion's `λ_y`-dependence
//!
//! ```text
//!   (n·p̃/2)·log((R_x + λ_y·R_y)/(n·p̃)) − (n·p_y/2)·log λ_y ,
//! ```
//!
//! (the second term is the `√λ_y` target-scaling Jacobian,
//! [`BehaviorBlock::reml_log_lambda_jacobian`]) whose unique stationary point is
//! the closed-form variance ratio
//!
//! ```text
//!   λ_y = (R_x/p_x) / (R_y/p_y)
//! ```
//!
//! ([`BehaviorBlock::reml_updated_log_lambda_y`]) — the classical REML estimate
//! of a variance-component ratio under a shared mean structure. The driver
//! below alternates (fit at fixed `λ_y`) ↔ (closed-form `λ_y` update at the
//! fitted residuals): block-coordinate descent on the joint profiled criterion,
//! each half-step solving its subproblem exactly. No grid search, no
//! user-tuned weight.
//!
//! # Payoffs realized here
//!
//! * **Gauge fixed by data** — the behavior block enters the same arrow-Schur
//!   inner solve that estimates `t`, so the latent coordinate is oriented by
//!   how the *output* changes, not by an arbitrary activation convention
//!   (see `tests_behavior_twoblock_rung2` for the planted case where the
//!   activation alone cannot orient `t` and the behavior block pins it).
//! * **Calibrated units** — the behavior target is nats-unit by construction
//!   ([`SphereTangentEmbedding`](crate::manifold::SphereTangentEmbedding)), and
//!   `λ_y` only sets the *inferential weight* of those units, so the fitted
//!   `C_k` always decodes to honest distributions ([`BehaviorBlock::split_decoder`]
//!   un-does the `√λ_y`).
//! * **Selection for mattering** — behaviorally inert structure has a zero
//!   behavior target; its residual variance ratio drives `λ_y` (and the atom's
//!   behavior decoder) toward earning nothing from the y-block evidence rather
//!   than manufacturing spurious weight.

use super::*;
use opt::{BacktrackConfig, backtracking_line_search};

/// Outcome of a two-block REML fit: the converged weight, the trajectory, and
/// the final inner loss.
#[derive(Clone, Debug)]
pub struct TwoBlockRemlFitReport {
    /// The REML-selected `log(λ_y)` installed on the term's behavior block.
    pub log_lambda_y: f64,
    /// Number of (fit, λ-update) outer sweeps performed (≥ 1).
    pub sweeps: usize,
    /// Whether the `log λ_y` fixed-point iteration met `log_lambda_tol` (as
    /// opposed to exhausting `max_sweeps`).
    pub converged: bool,
    /// `false` when the behavior residual carried no variance at some sweep
    /// (e.g. behavior constant across rows, target ≡ 0), in which case `λ_y`
    /// is not identifiable and was held at its last value — the honest report
    /// for a behaviorally inert block, not an error.
    pub lambda_identifiable: bool,
    /// Inner loss at the final fit (in the scaled augmented units of the last
    /// sweep's target).
    pub loss: SaeManifoldLoss,
    /// `log λ_y` after each sweep's update, for diagnostics (`sweeps` entries;
    /// the last equals `log_lambda_y` when identifiable).
    pub log_lambda_trajectory: Vec<f64>,
}

/// Caller-owned resolution knobs for [`SaeManifoldTerm::run_two_block_reml_fit`].
///
/// None of these choose *what* the fit converges to — the destination is the
/// data-determined REML stationary point. They only bound how far the inner
/// arrow-Schur solve and the outer `(fit, λ_y-update)` alternation walk toward
/// it, plus the pass-through inner regularization. Bundling them keeps the fit
/// entry point at a single grouped argument (the four inner knobs are the exact
/// pass-through set [`SaeManifoldTerm::run_joint_fit_arrow_schur`] already
/// takes, and the two outer knobs govern the sweep loop).
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct TwoBlockRemlControls {
    /// Upper bound on the outer `(fit, λ_y-update)` alternation. Must be ≥ 1.
    pub max_sweeps: usize,
    /// Inner arrow-Schur iteration cap, passed through to
    /// [`SaeManifoldTerm::run_joint_fit_arrow_schur`] unchanged each sweep.
    pub inner_max_iter: usize,
    /// Inner damped-Newton step size (finite, positive), passed through.
    pub step_size: f64,
    /// Ridge on the external-coordinate block of the inner solve, passed through.
    pub ridge_ext_coord: f64,
    /// Ridge on the decoder (`β`) block of the inner solve, passed through.
    pub ridge_beta: f64,
    /// Convergence tolerance on `|Δ log λ_y|` between sweeps (finite, positive).
    pub log_lambda_tol: f64,
}

/// Per-block outcome of a [`SaeManifoldTerm::run_multiblock_reml_fit`]: the
/// converged weight, whether it was identifiable, and its `log λ_ℓ` trajectory.
#[derive(Clone, Debug)]
pub struct BlockRemlOutcome {
    /// The block's label (carried through from its [`OutputBlock`]).
    pub label: String,
    /// The REML-selected `log(λ_ℓ)` installed on the block.
    pub log_lambda: f64,
    /// `false` when the block's residual carried no variance at some sweep (its
    /// target ≡ 0, or it was perfectly reconstructed): `λ_ℓ` is not identifiable
    /// and was held at its last value — the honest report for an inert block.
    pub identifiable: bool,
    /// `log λ_ℓ` after each sweep's update, for diagnostics.
    pub trajectory: Vec<f64>,
}

/// Outcome of a multi-block REML fit: the per-block converged weights, the sweep
/// count, and the final inner loss. The two-block [`TwoBlockRemlFitReport`] is
/// the `blocks.len() == 1` special case.
#[derive(Clone, Debug)]
pub struct MultiBlockRemlFitReport {
    /// One outcome per output block, in the order the blocks were supplied.
    pub blocks: Vec<BlockRemlOutcome>,
    /// Number of (fit, λ-update) outer sweeps performed (≥ 1).
    pub sweeps: usize,
    /// Whether the sweep loop stopped on convergence (all identifiable blocks
    /// met `log_lambda_tol`, or every block became unidentifiable) rather than
    /// exhausting `max_sweeps`.
    pub converged: bool,
    /// Inner loss at the final fit (in the scaled augmented units of the last
    /// sweep's targets).
    pub loss: SaeManifoldLoss,
}

impl SaeManifoldTerm {
    /// Run the block-generic multi-block joint fit with every block weight
    /// `λ_ℓ` selected by REML.
    ///
    /// The augmented target is `Z̃ = [Z | √λ_1·Y_1 | … | √λ_{K-1}·Y_{K-1}]`
    /// ([`stack_augmented_target`]): the anchor `Z` (`anchor`, `n × p_x`) plus
    /// the `√λ_ℓ`-scaled targets of `blocks`, all decoded from ONE shared latent
    /// coordinate and gate through the ordinary arrow-Schur joint fit. Each sweep
    /// (1) fits at the current weights, then (2) applies the closed-form REML
    /// variance-ratio update `λ_ℓ = (R_x/p_x)/(R_ℓ/p_ℓ)` to every block from the
    /// fitted residual — the *joint* stationary point, decoupled per block (see
    /// [`OutputBlock`]). The alternation is block-coordinate descent on the joint
    /// profiled criterion; `controls` bounds only how far it walks, never the
    /// destination.
    ///
    /// The term's atoms must be built at the augmented width
    /// `p̃ = p_x + Σ_ℓ p_ℓ`. On return the term holds the fitted state at the
    /// selected weights and each block in `blocks` carries its converged
    /// `log λ_ℓ`, so a decoder slice for block `ℓ` un-scaled by
    /// [`OutputBlock::split_honest_decoder`] decodes in honest units.
    ///
    /// A block whose residual carries no variance is *held* (its `λ_ℓ` frozen,
    /// `identifiable = false`) and excluded from the convergence test; when every
    /// block is held the sweep loop stops. At `blocks.len() == 1` this is
    /// bit-for-bit the two-block driver ([`Self::run_two_block_reml_fit`], which
    /// delegates here).
    pub fn run_multiblock_reml_fit(
        &mut self,
        anchor: ArrayView2<'_, f64>,
        blocks: &mut [OutputBlock],
        rho: &mut SaeManifoldRho,
        analytic_penalties: Option<&AnalyticPenaltyRegistry>,
        controls: TwoBlockRemlControls,
    ) -> Result<MultiBlockRemlFitReport, String> {
        let TwoBlockRemlControls {
            max_sweeps,
            inner_max_iter,
            step_size,
            ridge_ext_coord,
            ridge_beta,
            log_lambda_tol,
        } = controls;
        if max_sweeps == 0 {
            return Err(
                "SaeManifoldTerm::run_multiblock_reml_fit: max_sweeps must be ≥ 1".to_string(),
            );
        }
        if !(log_lambda_tol.is_finite() && log_lambda_tol > 0.0) {
            return Err(format!(
                "SaeManifoldTerm::run_multiblock_reml_fit: log_lambda_tol must be finite and \
                 positive; got {log_lambda_tol}"
            ));
        }
        if blocks.is_empty() {
            return Err(
                "SaeManifoldTerm::run_multiblock_reml_fit: need at least one output block"
                    .to_string(),
            );
        }
        let px = anchor.ncols();
        let p_tot = px + blocks.iter().map(|b| b.block_dim()).sum::<usize>();
        if p_tot != self.output_dim() {
            return Err(format!(
                "SaeManifoldTerm::run_multiblock_reml_fit: augmented width p_x + Σ p_ℓ = {p_tot} \
                 but the term's output_dim is {} (atoms must be built at the augmented width)",
                self.output_dim()
            ));
        }

        let n_obs = anchor.nrows();
        // Stacked-column bookkeeping owner. The block WIDTHS (hence every column
        // range) are fixed for the whole fit, so one layout serves every in-loop
        // range read (RSS spans, per-block decoder rescale); its stored `log λ_ℓ`
        // is the sweep-entry value and is NOT read for ranging. The FITTED layout
        // (converged weights) is installed on the term at the end.
        let range_layout = CrosscoderLayout::from_blocks(px, blocks);
        let dims: Vec<usize> = range_layout.block_dims().to_vec();
        let mut cur_log_lambda: Vec<f64> = blocks.iter().map(OutputBlock::log_lambda).collect();
        let mut trajectories: Vec<Vec<f64>> = vec![Vec::with_capacity(max_sweeps); blocks.len()];
        let mut identifiable = vec![true; blocks.len()];
        let mut converged = false;
        let mut sweeps = 0usize;

        // Backtracking bounds (resolution knobs, not answer-changers: the accepted
        // destination is a criterion stationary point either way — these only cap
        // how a single λ step is shortened when the full closed-form step would
        // INCREASE the profiled criterion, i.e. when the naive fixed-point would
        // overshoot into a block-abandonment runaway).
        const MAX_BACKTRACK: usize = 12;

        // Initial fit at the starting weights. Besides seeding the warm start and
        // `loss`, this runs BEFORE the first `fit_state_snapshot`, so any one-shot
        // atom rank reduction ([`Self::reduce_atoms_to_data_supported_rank`]) has
        // already settled the decoder/basis widths: every in-loop snapshot and its
        // in-place `.assign` restore then share one width and cannot shape-mismatch.
        let augmented = stack_augmented_target(anchor, blocks)?;
        let mut loss = self.run_joint_fit_arrow_schur(
            augmented.view(),
            rho,
            analytic_penalties,
            inner_max_iter,
            step_size,
            ridge_ext_coord,
            ridge_beta,
        )?;
        // No fitted RESIDUAL is carried across sweeps. Each sweep recomputes the
        // residual from a fresh refit-at-current-λ (the F1 baseline below), so the
        // closed-form λ* and the Armijo reference both read residuals produced with
        // the SAME inner budget the trials get; only `cur_log_lambda` and the
        // warm-start term state persist between sweeps. (`augmented` above seeded
        // the warm start and `loss`.)
        while sweeps < max_sweeps {
            sweeps += 1;

            // Sweep-entry warm-start state. Both the current-λ baseline and every
            // backtracking trial refit from THIS same snapshot with the SAME inner
            // budget, so their criteria differ only by the λ they were fit at.
            let base = self.fit_state_snapshot();

            // F1 — refit-at-current-λ baseline. A backtracking trial runs a full
            // `run_joint_fit_arrow_schur` with a fresh `inner_max_iter` budget, so
            // it can lower the profiled criterion purely by spending the truncated
            // inner solve's LEFTOVER descent at the SAME λ — inner-solve progress
            // masquerading as λ progress. To make the Armijo test a test of λ, the
            // value a trial must beat is itself a refit from `base` with the same
            // budget at the UNCHANGED λ: both sides start from one state and spend
            // one budget, differing only in λ. This baseline also gives the closed-
            // form λ* the residuals of the fit AT the current λ, not stale residuals
            // from the previous sweep's (differently-weighted) fit.
            self.fit_state_restore(&base)?;
            let base_aug = stack_augmented_target(anchor, blocks)?;
            let base_loss = self.run_joint_fit_arrow_schur(
                base_aug.view(),
                rho,
                analytic_penalties,
                inner_max_iter,
                step_size,
                ridge_ext_coord,
                ridge_beta,
            )?;
            let (base_rx, base_scaled) =
                self.augmented_block_rss(base_aug.view(), rho, &range_layout)?;
            let base_unscaled: Vec<f64> = (0..blocks.len())
                .map(|i| base_scaled[i] / blocks[i].lambda())
                .collect();
            // Envelope-priced penalty energy `P` of the CURRENT-λ refit (#2228):
            // twice the non-data-fit penalized-objective energy the inner engine
            // drove to (`P = 2·(penalized_objective_total − data_fit)` — decoder
            // smoothness + ARD + assignment prior + any analytic-registry energy).
            // Since `2e178664f` the inner engine returns the PENALIZED-objective
            // fixed point, so its residuals carry the penalty trade-off; pricing
            // `P` into the pooled dispersion is what restores the envelope theorem
            // (the profiled criterion's total λ-derivative equals its partial), so
            // the closed-form `λ*` is exact again. Held FIXED across a sweep's λ
            // moves (it is a fitted-state quantity, not an explicit λ function).
            let base_pen = 2.0
                * (self.penalized_objective_total(
                    base_aug.view(),
                    rho,
                    analytic_penalties,
                    1.0,
                )? - base_loss.data_fit);
            let baseline_crit = profiled_penalized_quasi_laplace_criterion(
                n_obs,
                px,
                base_rx,
                &base_unscaled,
                &dims,
                &cur_log_lambda,
                base_pen,
            )?;
            // The improved current-λ fit to fall back to if no λ step is accepted.
            let baseline_state = self.fit_state_snapshot();
            // The term currently holds this baseline fit; record its loss so any
            // convergence/stall break below reports the state actually installed.
            loss = base_loss;

            // Closed-form λ_ℓ target per identifiable block from the CURRENT-λ
            // refit's unscaled residuals (λ_ℓ* = (R_x/p_x)/(R_ℓ/p_ℓ)).
            let mut log_lambda_star = cur_log_lambda.clone();
            let mut any_update = false;
            let mut max_delta = 0.0_f64;
            for idx in 0..blocks.len() {
                if !identifiable[idx] {
                    continue;
                }
                let r_ell = base_unscaled[idx];
                // `(R_x + P)/p_x` — the penalty-priced anchor variance. The
                // coupled multiblock fixed point `λ_ℓ·R_ℓ = d_ℓ·pooled'/p̃`
                // collapses to `λ_ℓ* = ((R_x+P)/p_x)/(R_ℓ/d_ℓ)`, so `P` enters
                // the closed form only through this numerator (see
                // `profiled_penalized_quasi_laplace_block_efs_log_lambda_steps`).
                let var_x = (base_rx + base_pen) / px as f64;
                if !(r_ell > 0.0) || !(var_x > 0.0) {
                    // No block residual variance ⇒ λ_ℓ unidentifiable; hold it.
                    identifiable[idx] = false;
                    continue;
                }
                let var_y = r_ell / dims[idx] as f64;
                let star = (var_x / var_y).ln();
                log_lambda_star[idx] = star;
                max_delta = max_delta.max((star - cur_log_lambda[idx]).abs());
                any_update = true;
            }

            // Sufficient-decrease floor for the line search's Armijo test. It is
            // the ONLY authority that decides both a λ move and convergence below:
            // the closed-form λ* is used solely as a search DIRECTION, never as a
            // claimed criterion stationary point, so no `predicted_decrease` gate
            // that prices the closed form against HELD residuals is consulted (that
            // gate desynced from the refit — see the envelope note below).
            let armijo_eps = 1e-9 * (1.0 + baseline_crit.abs());

            if !any_update || max_delta <= log_lambda_tol {
                // Every identifiable block held (no weight can move) or the closed
                // form proposes no move beyond tolerance — the EFS fixed point. The
                // term holds the current-λ baseline fit, which is the answer.
                converged = true;
                for idx in 0..blocks.len() {
                    trajectories[idx].push(cur_log_lambda[idx]);
                }
                break;
            }

            // ENVELOPE-FREE CONVERGENCE (#2015). The prior fix priced a penalty
            // energy `P = 2·(penalized_objective − data_fit)` into the pooled
            // dispersion to restore the envelope theorem, so the closed-form λ*
            // (the per-block variance-ratio root at HELD residuals) would be the
            // profiled criterion's stationary point and `baseline_crit − C(λ*)`
            // a valid predicted decrease. That envelope is BROKEN: the sweep-first
            // inner engine (`converge_inner_for_undamped_logdet`) converges to the
            // penalized quasi-Laplace fixed point, whose stationarity condition
            // ∂(data_fit + penalty + logdet)/∂θ̂ = 0 leaves ∂(data_fit + penalty)/∂θ̂
            // = −∂logdet/∂θ̂ ≠ 0 — `penalized_objective_total` excludes the Laplace
            // log-determinant, so pricing `P` cancels only the penalty gradient,
            // not the logdet one. Hence `∂pooled'/∂θ̂ ≠ 0`, the profiled criterion's
            // TOTAL λ-derivative keeps non-vanishing dR_x/dλ, dR_ℓ/dλ, dP/dλ terms
            // (mediated by dθ̂/dλ), and the closed-form λ* is NOT the stationary
            // point of the criterion the trials actually evaluate (C at REFIT
            // residuals). Symptom: λ* proposed a `> tol` move whose held-residual
            // predicted decrease exceeded `armijo_eps`, yet every backtracking
            // fraction of the REFIT criterion failed Armijo — the sweep parked at a
            // bit-identical λ with `converged = false` (the planted-ratio failure).
            //
            // The fix is to let the EXACT refit criterion decide everything. The
            // closed form only orients the 1-D search direction `d = λ* − λ_cur`;
            // acceptance and convergence are read off `C(refit(λ))`, guaranteeing
            // consistency by construction. We backtrack along `+d` AND, if no
            // forward fraction descends, along `−d`: if neither direction lowers
            // the refit criterion below `baseline_crit − armijo_eps`, then λ_cur is
            // a local minimiser of the exact criterion along the closed-form axis
            // to the search resolution — a GENUINE fixed point, so `converged =
            // true`. Only a realizable descent the search actually captured (in
            // either direction) keeps the sweep moving.

            // One signed-step trial of the closed-form direction `d = λ* − λ_cur`.
            // `s` is a SIGNED scale: `s > 0` steps toward λ*, `s < 0` steps away
            // from it (the reverse half-axis). Each trial refits from the SAME
            // `base` snapshot with the SAME inner budget as the baseline, so the
            // comparison isolates the λ move; the criterion penalises λ→0, so a step
            // that would abandon a block is rejected — the fix that makes the
            // otherwise non-contractive (fit, λ-update) alternation monotone. The
            // trial threads `(trial_ll, trial_loss)` through the payload.
            let accepted_step = {
                let mut trial =
                    |s: f64| -> Result<Option<(f64, (Vec<f64>, SaeManifoldLoss))>, String> {
                        let mut trial_ll = cur_log_lambda.clone();
                        for idx in 0..blocks.len() {
                            if identifiable[idx] {
                                trial_ll[idx] = cur_log_lambda[idx]
                                    + s * (log_lambda_star[idx] - cur_log_lambda[idx]);
                            }
                        }
                        for idx in 0..blocks.len() {
                            blocks[idx] = blocks[idx].with_log_lambda(trial_ll[idx])?;
                        }
                        self.fit_state_restore(&base)?;
                        // EXACT warm-start reparameterization: the stacked target's block
                        // columns are `√λ_ℓ·Y_ℓ`, and quadratic β-penalties are scale-
                        // equivariant (`min_β ‖√λ·y − Φβ‖² + βᵀSβ` maps to
                        // `λ(‖y − Φb‖² + bᵀSb)` under `β = √λ·b`), so the incumbent fit at
                        // `λ_cur` maps to the EXACT incumbent at `λ_trial` by scaling each
                        // behavior block's decoder columns by `√(λ_trial/λ_cur)` (valid
                        // for a step of EITHER sign). Without this rescale every trial
                        // starts with a spurious scaled-residual `(√λ_t − √λ_c)²·‖Ŷ‖²`
                        // that the TRUNCATED inner solve must burn its iteration budget
                        // removing — which biases the comparison against exactly the
                        // large closed-form λ jumps the fixed point needs from a distant
                        // start (the from-0.0 20-sweep-exhaustion failure of
                        // `reml_selects_lambda_y…`). The rescaled state is the same fit in
                        // the new parameterization, so this changes no fixed point — it
                        // only stops the search from paying an artificial re-fitting tax
                        // proportional to the step.
                        for idx in 0..blocks.len() {
                            let scale = gam_problem::checked_exp_log_strength(
                                0.5 * (trial_ll[idx] - cur_log_lambda[idx]),
                            )
                            .map_err(|error| {
                                format!(
                                    "behavior smoothing warm-start rescale is outside the canonical log-strength domain: {error}"
                                )
                            })?;
                            if scale != 1.0 {
                                let range = range_layout.block_range(idx);
                                for atom in &mut self.atoms {
                                    let mut cols =
                                        atom.decoder_coefficients_mut().slice_move(s![.., range.clone()]);
                                    cols.mapv_inplace(|v| v * scale);
                                }
                            }
                        }
                        let aug = stack_augmented_target(anchor, blocks)?;
                        let trial_loss = self.run_joint_fit_arrow_schur(
                            aug.view(),
                            rho,
                            analytic_penalties,
                            inner_max_iter,
                            step_size,
                            ridge_ext_coord,
                            ridge_beta,
                        )?;
                        let (trx, trb_scaled) =
                            self.augmented_block_rss(aug.view(), rho, &range_layout)?;
                        let trb_unscaled: Vec<f64> = (0..blocks.len())
                            .map(|i| trb_scaled[i] / trial_ll[i].exp())
                            .collect();
                        // The trial's OWN realized penalty energy `P` at its λ-trial
                        // fitted state (same `2·(penalized_objective_total − data_fit)`
                        // as the baseline). This is NOT an envelope correction (the
                        // envelope is broken; see the note above) — it simply makes both
                        // sides of the comparison price their fitted-state penalty the
                        // same way, so `trial_crit` and `baseline_crit` are the SAME
                        // profiled criterion read at two λ's.
                        let trial_pen = 2.0
                            * (self.penalized_objective_total(
                                aug.view(),
                                rho,
                                analytic_penalties,
                                1.0,
                            )? - trial_loss.data_fit);
                        let trial_crit = profiled_penalized_quasi_laplace_criterion(
                            n_obs,
                            px,
                            trx,
                            &trb_unscaled,
                            &dims,
                            &trial_ll,
                            trial_pen,
                        )?;
                        Ok(Some((trial_crit, (trial_ll, trial_loss))))
                    };

                // Backtrack along `+d` toward λ* (initial step 1.0, ×0.5 contraction);
                // if no forward fraction descends the refit criterion, backtrack along
                // `−d` away from λ* (initial step −1.0, contraction preserves the sign).
                // The same trial closure serves both directions — `s` carries the sign.
                // `accept` is the exact sufficient-decrease test
                // `trial_crit < baseline_crit − armijo_eps`, direction-agnostic.
                let forward = backtracking_line_search::<(Vec<f64>, SaeManifoldLoss), String>(
                    BacktrackConfig {
                        initial_step: 1.0,
                        contraction: 0.5,
                        max_steps: MAX_BACKTRACK,
                    },
                    &mut trial,
                    |_, trial_crit| trial_crit < baseline_crit - armijo_eps,
                )?;
                if forward.is_some() {
                    forward
                } else {
                    backtracking_line_search::<(Vec<f64>, SaeManifoldLoss), String>(
                        BacktrackConfig {
                            initial_step: -1.0,
                            contraction: 0.5,
                            max_steps: MAX_BACKTRACK,
                        },
                        &mut trial,
                        |_, trial_crit| trial_crit < baseline_crit - armijo_eps,
                    )?
                }
            };

            if let Some(step) = accepted_step {
                let (trial_ll, trial_loss) = step.payload;
                cur_log_lambda = trial_ll;
                loss = trial_loss;
            } else {
                // Neither `+d` nor `−d` lowers the EXACT refit criterion below
                // `baseline_crit − armijo_eps`: λ_cur is a local minimiser of the
                // criterion the trials evaluate, along the closed-form axis, to the
                // search resolution — a GENUINE criterion fixed point. Reinstate the
                // improved current-λ fit (`baseline_state`, no worse than the
                // sweep-entry state) at the unchanged weights and report converged.
                //
                // This is the #2015 fix: the old code reached here only via the
                // forward search and declared a STALL (`converged = false`),
                // presuming the closed-form `+d` was a true descent direction the
                // truncated solve just could not realize. With the envelope broken
                // (see the note above) `+d` is NOT the refit criterion's descent
                // direction, so a forward-only failure was misread as a stall for a
                // λ that had actually settled. Testing BOTH directions on the exact
                // criterion turns "no move helps" into an honest convergence verdict
                // (and still catches a real reverse-descent, taking that step above).
                for idx in 0..blocks.len() {
                    blocks[idx] = blocks[idx].with_log_lambda(cur_log_lambda[idx])?;
                }
                self.fit_state_restore(&baseline_state)?;
                loss = base_loss;
                converged = true;
                for idx in 0..blocks.len() {
                    trajectories[idx].push(cur_log_lambda[idx]);
                }
                break;
            }
            for idx in 0..blocks.len() {
                trajectories[idx].push(cur_log_lambda[idx]);
            }
        }
        // Install the FITTED stacked-column layout (converged per-block weights)
        // so `layer_decoder(k, ℓ)` returns each layer's honest-units decoder with
        // no caller re-slicing. `blocks` now carry the converged `log λ_ℓ`; the
        // width equals `output_dim()` (validated above), so the install succeeds.
        self.set_crosscoder_layout(CrosscoderLayout::from_blocks(px, blocks))?;

        let outcomes = blocks
            .iter()
            .enumerate()
            .map(|(idx, block)| BlockRemlOutcome {
                label: block.label.clone(),
                log_lambda: block.log_lambda(),
                identifiable: identifiable[idx],
                trajectory: std::mem::take(&mut trajectories[idx]),
            })
            .collect();
        Ok(MultiBlockRemlFitReport {
            blocks: outcomes,
            sweeps,
            converged,
            loss,
        })
    }

    /// Anchor RSS `R_x` (over the `p_x` anchor columns) and the SCALED per-block
    /// RSS (over each block's column span) of the current fitted residual. The
    /// column offsets come from `layout` (the single owner of the stacked-column
    /// bookkeeping), not a by-hand `off_ℓ` accumulation.
    fn augmented_block_rss(
        &self,
        augmented: ArrayView2<'_, f64>,
        rho: &SaeManifoldRho,
        layout: &CrosscoderLayout,
    ) -> Result<(f64, Vec<f64>), String> {
        let residual = self.reconstruction_residual(augmented, rho)?;
        let px = layout.anchor_dim();
        let mut rss_x = 0.0_f64;
        for row in residual.rows() {
            for j in 0..px {
                let r = row[j];
                rss_x += r * r;
            }
        }
        let mut out = Vec::with_capacity(layout.num_blocks());
        for l in 0..layout.num_blocks() {
            let mut rss = 0.0_f64;
            for row in residual.rows() {
                for j in layout.block_range(l) {
                    let r = row[j];
                    rss += r * r;
                }
            }
            out.push(rss);
        }
        Ok((rss_x, out))
    }

    /// Install the crosscoder stacked-column layout on the term, validating that
    /// its total width matches the atoms' output dimension
    /// (`p_x + Σ_ℓ p_ℓ == output_dim()`). Called by
    /// [`Self::run_multiblock_reml_fit`] with the fitted per-block weights; a
    /// caller may also install one explicitly to enable [`Self::layer_decoder`]
    /// on a term whose decoders were fit elsewhere at the augmented width.
    pub fn set_crosscoder_layout(&mut self, layout: CrosscoderLayout) -> Result<(), String> {
        if layout.total_dim() != self.output_dim() {
            return Err(format!(
                "SaeManifoldTerm::set_crosscoder_layout: layout total width p_x + Σ p_ℓ = {} but \
                 the term's output_dim is {} (atoms must be built at the augmented width)",
                layout.total_dim(),
                self.output_dim()
            ));
        }
        self.crosscoder_layout = Some(layout);
        Ok(())
    }

    /// The installed crosscoder stacked-column layout, or `None` for a term with
    /// no multi-block layout recorded.
    pub fn crosscoder_layout(&self) -> Option<&CrosscoderLayout> {
        self.crosscoder_layout.as_ref()
    }

    /// The HONEST-units per-layer decoder `B_k^(ℓ)` of atom `k`, output block `ℓ`:
    /// the decoder columns of block `ℓ` divided by `√λ_ℓ` (un-doing the target
    /// scaling that `stack_augmented_target` applied), and — when a Tier-0
    /// column-equilibration scale is installed (#2015; see
    /// [`Self::set_tier0_scale`] and `crosscoder_fit::equilibrate_crosscoder_columns`)
    /// — un-doing that per-column scale too, so the returned decoder is honest in
    /// the CALLER's raw target units regardless of the internal conditioning
    /// frame the inner solve actually ran on. Requires an installed
    /// [`CrosscoderLayout`] (via a multi-block fit or [`Self::set_crosscoder_layout`]).
    ///
    /// This is the first-class form of the by-hand slice+unscale
    /// (`decoder_coefficients[:, off_ℓ..off_ℓ+p_ℓ] / √λ_ℓ`, times the column's
    /// tier0 scale): identical values, but the offsets, `√λ_ℓ`, and tier0 scale
    /// are owned here so no caller recomputes them.
    pub fn layer_decoder(&self, k: usize, l: usize) -> Result<Array2<f64>, String> {
        let layout = self.crosscoder_layout.as_ref().ok_or_else(|| {
            "SaeManifoldTerm::layer_decoder: no crosscoder layout installed (run \
             run_multiblock_reml_fit or set_crosscoder_layout first)"
                .to_string()
        })?;
        if k >= self.atoms.len() {
            return Err(format!(
                "SaeManifoldTerm::layer_decoder: atom index k={k} out of range (K = {})",
                self.atoms.len()
            ));
        }
        if l >= layout.num_blocks() {
            return Err(format!(
                "SaeManifoldTerm::layer_decoder: block index ℓ={l} out of range (L-1 = {})",
                layout.num_blocks()
            ));
        }
        let inv = 1.0 / layout.sqrt_lambda(l);
        // Materialize the full-width decoder (Tier-0 column scale already
        // undone) before crossing the layer boundary; the fit may use a
        // reduced Grassmann coordinate internally.
        let physical = self.tier0_unscaled_full_width_decoder(k);
        let scaled = physical.slice(s![.., layout.block_range(l)]);
        Ok(scaled.mapv(|value| inv * value))
    }

    /// The full-width (augmented) decoder of atom `k`, with the Tier-0
    /// column-equilibration scale (#2015; [`Self::set_tier0_scale`],
    /// `crosscoder_fit::equilibrate_crosscoder_columns`) undone column-by-column
    /// when one is installed — a strict no-op on the historical (unequilibrated)
    /// path. Every consumer that carves an honest per-layer decoder out of the
    /// full augmented width ([`Self::layer_decoder`] above, the crosscoder
    /// drift/transport reports) starts from this, so an equilibration scale
    /// installed by the crosscoder fit entry is undone exactly once, in one
    /// place, rather than re-derived at each call site.
    pub(crate) fn tier0_unscaled_full_width_decoder(&self, k: usize) -> Array2<f64> {
        let mut decoder = self.atoms[k].full_width_decoder();
        if let Some(scale) = self.tier0_scale() {
            for (col, &s) in scale.iter().enumerate() {
                decoder.column_mut(col).mapv_inplace(|v| v * s);
            }
        }
        decoder
    }

    /// Snapshot the ENTIRE mutable fit state a λ line-search trial perturbs, so a
    /// rejected trial rolls back to a bit-identical base.
    ///
    /// A trial runs a full [`Self::run_joint_fit_arrow_schur`], which moves far
    /// more than the decoder β and latent coords: it refreshes each atom's cached
    /// `basis_values` / `basis_jacobian` while preserving or structurally
    /// transporting the fixed reference `smooth_penalty`,
    /// rewrites the assignment `logits`, re-derives the active-set `last_row_layout`,
    /// and advances the Gumbel `temperature_schedule` one anneal step per inner
    /// iteration. The canonical [`Self::snapshot_mutable_state`] captures the first
    /// group (the same state the inner damped-Newton line search itself rolls back);
    /// the schedule is a per-call *stateful* counter, so a rejected trial that
    /// leaves it advanced would start the next trial at a colder temperature and
    /// make the backtracking trials incomparable — capture it here too. (An atom
    /// rank reduction, [`Self::reduce_atoms_to_data_supported_rank`], is idempotent
    /// after the sweep-entry fit, so the restored decoder width always matches.)
    fn fit_state_snapshot(&self) -> (SaeManifoldMutableState, Option<GumbelTemperatureSchedule>) {
        (
            self.snapshot_mutable_state(),
            self.temperature_schedule.clone(),
        )
    }

    /// Restore a [`Self::fit_state_snapshot`]. `restore_mutable_state` rebuilds
    /// each atom's `basis_values` / `basis_jacobian` from the restored coordinates
    /// (the differential snapshot stores only the cheap driving state), so the
    /// restored state is immediately consistent for a residual/fitted read.
    fn fit_state_restore(
        &mut self,
        snap: &(SaeManifoldMutableState, Option<GumbelTemperatureSchedule>),
    ) -> Result<(), String> {
        self.restore_mutable_state(&snap.0)?;
        self.temperature_schedule.clone_from(&snap.1);
        Ok(())
    }
}

impl SaeManifoldTerm {
    /// Run the Rung-2 two-block joint fit with `λ_y` selected by REML.
    ///
    /// Requires a [`BehaviorBlock`] installed via
    /// [`Self::set_behavior_block`]; `activation` is the raw activation target
    /// `Z` (`n × p_x`) — the augmented target is stacked internally at each
    /// sweep's current `λ_y`. `rho`, `analytic_penalties`, and the inner knobs
    /// carried in `controls` (`inner_max_iter`, `step_size`, and the two ridges)
    /// are passed through to [`Self::run_joint_fit_arrow_schur`] unchanged.
    ///
    /// `controls.max_sweeps` bounds the outer (fit, λ-update) alternation;
    /// `controls.log_lambda_tol` is the convergence tolerance on `|Δ log λ_y|`
    /// (see [`TwoBlockRemlControls`]: all of these are caller-owned resolution
    /// choices — not fit hyperparameters. The *destination* is the
    /// data-determined REML stationary point; these only bound how long we walk
    /// toward it).
    ///
    /// On return the term holds the fitted two-block state at the selected
    /// `λ_y` (its behavior block updated in place), so
    /// [`BehaviorBlock::split_decoder`] on the fitted decoders yields the
    /// activation decoder `B_k` and the nats-unit behavior decoder `C_k`.
    pub fn run_two_block_reml_fit(
        &mut self,
        activation: ArrayView2<'_, f64>,
        rho: &mut SaeManifoldRho,
        analytic_penalties: Option<&AnalyticPenaltyRegistry>,
        controls: TwoBlockRemlControls,
    ) -> Result<TwoBlockRemlFitReport, String> {
        // Two-block-specific validation and error messages, preserved verbatim
        // for existing callers; the fit itself is the `K = 2` (single output
        // block) special case of the block-generic driver, which it delegates to
        // so the two paths are bit-for-bit identical by construction.
        if controls.max_sweeps == 0 {
            return Err(
                "SaeManifoldTerm::run_two_block_reml_fit: max_sweeps must be ≥ 1".to_string(),
            );
        }
        if !(controls.log_lambda_tol.is_finite() && controls.log_lambda_tol > 0.0) {
            return Err(format!(
                "SaeManifoldTerm::run_two_block_reml_fit: log_lambda_tol must be finite and \
                 positive; got {}",
                controls.log_lambda_tol
            ));
        }
        let Some(block) = self.behavior_block().cloned() else {
            return Err(
                "SaeManifoldTerm::run_two_block_reml_fit: no behavior block installed \
                 (call set_behavior_block first)"
                    .to_string(),
            );
        };
        if activation.ncols() != block.activation_dim {
            return Err(format!(
                "SaeManifoldTerm::run_two_block_reml_fit: activation has {} columns; behavior \
                 block declares p_x = {}",
                activation.ncols(),
                block.activation_dim
            ));
        }

        // The behavior block IS an output block: its (unscaled) tangent target,
        // its width p_y, and its weight λ_y drive the same augmented fit and the
        // same variance-ratio update.
        let mut output_blocks = vec![OutputBlock::new(
            "behavior",
            block.target.clone(),
            block.log_lambda_y(),
        )?];
        let report = self.run_multiblock_reml_fit(
            activation,
            &mut output_blocks,
            rho,
            analytic_penalties,
            controls,
        )?;

        // Keep the term's installed behavior block in sync with the selected
        // weight — the two-block contract that `behavior_block()` (and hence
        // `split_decoder` / `augmented_target`) reflects the converged λ_y.
        let fitted_block = block.with_log_lambda_y(output_blocks[0].log_lambda())?;
        self.set_behavior_block(fitted_block)?;

        let outcome = &report.blocks[0];
        Ok(TwoBlockRemlFitReport {
            log_lambda_y: outcome.log_lambda,
            sweeps: report.sweeps,
            converged: report.converged,
            lambda_identifiable: outcome.identifiable,
            loss: report.loss,
            log_lambda_trajectory: outcome.trajectory.clone(),
        })
    }
}