gam-terms 0.3.152

use super::*;
use gam_linalg::utils::SPECTRAL_DEFLATION_REL_FLOOR;
use gam_problem::{LOG_STRENGTH_MAX, LOG_STRENGTH_MIN, checked_exp_log_strength};

/// Exact floating-point continuation of `log(p) + 1` on the support of a
/// representable softmax row. An underflowed probability is exactly zero in the
/// value path, so its entropy contribution and all local derivatives are zero;
/// using the same branch everywhere keeps value/gradient/Hessian consistent.
#[inline]
fn entropy_log_plus_one(p: f64) -> f64 {
    if p > 0.0 { p.ln() + 1.0 } else { 0.0 }
}

/// Smooth upper envelope of `|x|` — the soft-abs (pseudo-Huber) magnitude
///
/// ```text
///   σ_ε(x) = sqrt(x² + ε²)
/// ```
///
/// used in place of `|·|` wherever a Gershgorin radius `Σ_j|H_kj|` is
/// differentiated (#2339). Takes the SQUARED smoothing scale `eps_sq = ε²`
/// because every caller derives it as `ε₀²·‖H_k·‖₂²` and never needs `ε` itself;
/// passing the square also keeps the degenerate `‖H_k·‖₂² = 0` row on the exact
/// `σ_0 = |·|` branch instead of routing it through a `sqrt` that would have to
/// be undone.
///
/// **Majorization (`σ_ε(x) ≥ |x|`) is the whole point** and is a hard guarantee,
/// not an asymptotic one: the Gershgorin diagonal `D` is a Loewner majorizer of
/// the indefinite entropy Hessian ONLY because each term dominates `|H_kj|`, so a
/// smoothing that dips below `|x|` — `x·tanh(x/ε)` and `ε·ln cosh(x/ε)` both do —
/// silently invalidates `D ⪰ H` and lets the assembled evidence block go
/// indefinite. In exact arithmetic `sqrt(x² + ε²) ≥ sqrt(x²) = |x|`; in `f64`
/// the sum can round BELOW `x²` when `ε² < ulp(x²)/2` and land up to ~1.5 ulp
/// under `|x|`, so the rounding direction is pinned with `max(·, |x|)`. That max
/// binds only where its two arguments agree to within one ulp, so it is a
/// rounding-direction guard on an identity — not a clamp on a signed quantity,
/// and not a reintroduced kink.
///
/// Gap: `0 ≤ σ_ε(x) − |x| = ε²/(σ_ε(x) + |x|) ≤ ε`, attained at `x = 0` where
/// `σ_ε(0) = ε` — the envelope is strictly above `|·|` exactly at the seam it
/// exists to smooth, and collapses onto `|·|` like `ε²/(2|x|)` away from it.
#[inline]
#[must_use]
pub fn soft_abs_squared_scale(x: f64, eps_sq: f64) -> f64 {
    (x * x + eps_sq).sqrt().max(x.abs())
}

// ---------------------------------------------------------------------------
// Sparsity penalty
// ---------------------------------------------------------------------------

/// Sparsifier kernel.
///
/// * `SmoothedL1 { eps }` — `Σ_i sqrt(x_i² + ε²)`. The smoothing scale `ε`
///   may be REML-selected, in which case the
///   shrink rate `ε → 0` is governed by the marginal likelihood (Occam keeps
///   `ε` large when the data don't demand sharpness).
/// * `Hoyer` — `(√n · ‖x‖_1 − ‖x‖_2) / (√n − 1)`. Scale-invariant; encourages
///   absolute sparsity even when the global scale of `x` drifts.
/// * `Log { delta }` — `Σ_i log(1 + x_i² / δ²)`. Strongly concave; aggressive
///   sparsifier suitable for active-set / iterative-reweighted paths.
#[derive(Debug, Clone, Copy)]
pub enum SparsityKind {
    SmoothedL1 { eps: f64 },
    Hoyer,
    Log { delta: f64 },
}

/// Sparsity penalty on a slice of β (SAE codes) or ext-coords (soft atom assignments).
///
/// The smoothed-L¹ default `Σ_i sqrt(x_i² + ε²)` is the simplest analytic
/// option. Its gradient is `x_i / sqrt(x_i² + ε²)` (a smooth sign function),
/// and its Hessian is diagonal with entries `ε² / (x_i² + ε²)^{3/2}` — so
/// `hvp` is cheap and the inner Newton step inherits a benign block-diagonal
/// regularizer.
///
/// When to use: any time a parameter block carries a "this should be sparse"
/// prior — SAE atom codes (β slice), soft-routing weights on a latent
/// ext-coordinate slice. For SAE codes specifically, smoothed-L¹ with REML-selected `ε`
/// gives the principled relaxation of the L¹ objective without giving up
/// differentiability.
#[derive(Debug, Clone)]
pub struct SparsityPenalty {
    pub target_tier: PenaltyTier,
    pub kind: SparsityKind,
    pub weight: f64,
    pub weight_schedule: Option<ScalarWeightSchedule>,
    /// Whether local rho coordinate 1 learns `log ε` (or `log δ`). Coordinate
    /// 0 is always the log-strength. Keeping this as a boolean makes invalid
    /// local index layouts unrepresentable.
    learnable_smoothing: bool,
}

/// Entropy sparsity over row-wise softmax assignment logits.
///
/// This is the SAE-manifold soft-assignment penalty. The target is a flat
/// row-major `(N, K)` logit matrix. Assignments are
/// `a_i = softmax(logits_i / temperature)`, and the penalty is
///
/// ```text
///   lambda_sparse * sum_i H(a_i)
///   H(a_i) = -sum_k a_ik log a_ik
/// ```
///
/// Minimizing entropy drives each row toward a small active support while the
/// softmax keeps `a_ik >= 0` and `sum_k a_ik = 1`. The exact Hessian is dense
/// in each row and can be indefinite because entropy is concave in assignment
/// space, so callers must use the HVP rather than a diagonal Hessian shortcut.
#[derive(Debug, Clone)]
pub struct SoftmaxAssignmentSparsityPenalty {
    pub k_atoms: usize,
    pub temperature: f64,
    pub weight: f64,
    pub weight_schedule: Option<ScalarWeightSchedule>,
    /// #991 design-honesty per-row weights `w_i` (mean-1). When present, row `i`'s
    /// prior contribution is scaled by `w_i` in EVERY aggregate channel — value,
    /// `grad_target`, `hessian_diag`, `hvp`, `psd_majorizer_diag`, `grad_rho`.
    /// Because each of those is linear in the per-row penalty strength, scaling
    /// the strength by `w_i` scales all channels by the same `w_i` and cannot
    /// desync them (the value/gradient FD oracle gates this). The per-row *block*
    /// helpers (`row_dense_hessian` / `row_psd_majorizer` / their logit
    /// derivatives / `psd_majorizer_abs_row_sums`) take an explicit `scale` and a
    /// single row, so their callers apply `scale·w_i` instead. `None` ⇒ every
    /// weight is `1`, bit-for-bit the unweighted path.
    pub row_weights: Option<std::sync::Arc<[f64]>>,
}

impl SoftmaxAssignmentSparsityPenalty {
    #[must_use]
    pub fn new(k_atoms: usize, temperature: f64) -> Self {
        assert!(k_atoms > 0);
        assert!(temperature > 0.0);
        Self {
            k_atoms,
            temperature,
            weight: 1.0,
            weight_schedule: None,
            row_weights: None,
        }
    }

    /// Install #991 design-honesty per-row weights (see [`Self::row_weights`]).
    /// A uniform / absent design is passed as `None` so the unweighted arithmetic
    /// stays bit-for-bit; a present slice must have one finite weight per row.
    #[must_use]
    pub fn with_row_weights(mut self, weights: Option<&[f64]>) -> Self {
        self.row_weights = weights.map(|w| std::sync::Arc::from(w.to_vec()));
        self
    }

    /// Per-row strength multiplier `w_i` (defaults to `1.0` when no design weights
    /// are installed). Callers of the per-row *block* helpers fold this into the
    /// `scale` they pass so those channels carry the identical weighting.
    #[must_use]
    pub fn row_weight(&self, row: usize) -> f64 {
        self.row_weights.as_ref().map_or(1.0, |w| w[row])
    }

    impl_with_weight_schedule!(weight);

    fn softmax_row(&self, row: &[f64]) -> Vec<f64> {
        let inv_tau = 1.0 / self.temperature;
        let mut max_logit = f64::NEG_INFINITY;
        for (idx, &v) in row.iter().enumerate() {
            assert!(
                v.is_finite(),
                "SoftmaxAssignmentSparsityPenalty: non-finite logit at atom {idx}: {v}"
            );
            max_logit = max_logit.max(v);
        }
        let mut out = vec![0.0; self.k_atoms];
        let mut sum = 0.0;
        for i in 0..self.k_atoms {
            let v = ((row[i] - max_logit) * inv_tau).exp();
            out[i] = v;
            sum += v;
        }
        assert!(
            sum.is_finite() && sum > 0.0,
            "SoftmaxAssignmentSparsityPenalty: non-finite softmax normalizer"
        );
        for v in out.iter_mut() {
            *v /= sum;
        }
        out
    }

    /// Dimensionless soft-abs temperature `ε₀` for the smooth Gershgorin
    /// majorizer (#2339). NOT a tunable knob — derived below, and derived from
    /// the problem's own dictionary size `k_atoms`.
    ///
    /// [`Self::psd_majorizer_abs_row_sums`] smooths the DIMENSIONLESS normalized
    /// row entries `u_kj = H_kj/‖H_k·‖₂` (which satisfy `Σ_j u_kj² = 1`, hence
    /// `|u_kj| ≤ 1` — a quantity whose natural scale is exactly unity, the direct
    /// analogue of the ARD half's `cos κt`) and multiplies back by the row's own
    /// curvature scale `‖H_k·‖₂`. The per-entry envelope gap is at most `ε₀` in
    /// those units, so over the `K` terms of a row sum
    ///
    /// ```text
    ///   0 ≤ D̃_kk − D_kk ≤ K·ε₀·‖H_k·‖₂ ≤ K·ε₀·D_kk        (‖·‖₂ ≤ ‖·‖₁ = D_kk)
    /// ```
    ///
    /// — a purely RELATIVE gap. The criterion resolves relative curvature only
    /// down to the spectral-deflation floor [`SPECTRAL_DEFLATION_REL_FLOOR`]
    /// (`λ < floor·λ_max` is deflated as null), so requiring the majorization gap
    /// to sit at that floor,
    ///
    /// ```text
    ///   K·ε₀ ≤ SPECTRAL_DEFLATION_REL_FLOOR
    /// ```
    ///
    /// and taking the binding (largest-admissible, hence smoothest) value gives
    ///
    /// ```text
    ///   ε₀ = SPECTRAL_DEFLATION_REL_FLOOR / K.
    /// ```
    ///
    /// The absolute smoothing scale actually applied, `ε_k = ε₀·‖H_k·‖₂`, is read
    /// entirely off the row's own curvature; the only constant involved is the
    /// floor the rest of the engine already resolves against. This is the same
    /// statement the ARD half proves for its softplus clamp
    /// (`α·τ₀·ln2 = α·floor`).
    #[must_use]
    pub fn soft_abs_temperature(k_atoms: usize) -> f64 {
        SPECTRAL_DEFLATION_REL_FLOOR / (k_atoms as f64)
    }

    /// Smoothed absolute row sums of the exact per-row dense entropy Hessian,
    /// used as a Gershgorin / diagonal-dominance PSD majorizer.
    ///
    /// The exact per-row Hessian wrt logits (symmetric, dense) is
    ///
    /// ```text
    ///   H_kj = (λ/τ²)·a_k·[ δ_kj·(m − L_k − 1) + a_j·(L_k + L_j + 1 − 2m) ],
    ///   L_k = ln a_k + 1,   m = Σ_j a_j L_j,
    /// ```
    ///
    /// whose diagonal coincides with [`AnalyticPenalty::hessian_diag`]. Entropy
    /// is concave in assignment space, so this block is indefinite (negative on
    /// near-uniform rows). Setting `D_kk = Σ_j |H_kj|` makes `D − H` symmetric
    /// with nonnegative diagonal and diagonally dominant
    /// (`D_kk − H_kk = |H_kk| − H_kk + Σ_{j≠k}|H_kj| ≥ Σ_{j≠k}|(D−H)_kj|`),
    /// hence PSD: `D ⪰ H` and `D ⪰ 0` both hold. `D` is a genuine PSD diagonal
    /// operator that dominates the dense Hessian's quadratic form — unlike the
    /// raw indefinite diagonal, which is neither PSD nor a faithful stand-in for
    /// the dense operator.
    ///
    /// # The `|·|` is smoothed (#2339)
    ///
    /// Every off-diagonal `H_kj` crosses zero on the codimension-1 surface
    /// `L_k + L_j + 1 = 2m`, so the raw radius `Σ_j|H_kj|` carries a kink through
    /// generic logit space and its θ-adjoint `Σ_j sign(H_kj)·Ḣ_kj` JUMPS across
    /// it — the objective↔gradient desync that stalls any outer method
    /// differentiating the streaming criterion `½log|B̃|`. This returns instead
    ///
    /// ```text
    ///   D̃_kk = Σ_j σ_{ε_k}(H_kj) = Σ_j sqrt(H_kj² + ε₀²·‖H_k·‖₂²),
    /// ```
    ///
    /// the soft-abs envelope [`soft_abs_squared_scale`] applied at the row's own
    /// scale, with `ε₀` from [`Self::soft_abs_temperature`]. Four properties, all
    /// gated by `soft_abs_gershgorin_2339_tests`:
    ///
    /// 1. **Majorizer.** `σ_ε(x) ≥ |x|` entrywise ⇒ `D̃_kk ≥ D_kk ≥ 0`, so
    ///    `D̃ − D` is a nonnegative diagonal and `D̃ − H = (D̃ − D) + (D − H)` is a
    ///    sum of two PSD matrices. `D̃ ⪰ D ⪰ H` and `D̃ ⪰ D ⪰ 0` are INHERITED,
    ///    never re-argued: smoothing can only move the bound in the safe
    ///    direction.
    /// 2. **Smooth.** `H_kj² + ε₀²Σ_l H_kl²` is a polynomial in the (analytic)
    ///    entries of `H_k·` and is `≥ ε₀²‖H_k·‖₂² > 0` whenever the row is
    ///    nonzero, and `sqrt` is analytic on `(0,∞)`, so `D̃_kk` is real-analytic
    ///    (C^ω) wherever `H_k· ≠ 0` — in particular across every individual
    ///    zero crossing. The only surviving non-smooth point is the SIMULTANEOUS
    ///    vanishing `H_k· = 0` (codimension `K`), which here happens exactly when
    ///    `a_k` underflows to 0; there `H_k· ≡ 0` and `Ḣ_k· ≡ 0`, so value and
    ///    derivative are identically zero and the exact-zero continuation that
    ///    `entropy_log_plus_one` already uses carries through.
    /// 3. **Tight.** `0 ≤ D̃_kk − D_kk ≤ K·ε₀·‖H_k·‖₂ ≤ SPECTRAL_DEFLATION_REL_FLOOR·D_kk`
    ///    (derivation in [`Self::soft_abs_temperature`]) — below the relative
    ///    resolution at which the factorization declares a direction null.
    /// 4. **Scale-derived.** `σ` is applied at `ε_k = ε₀‖H_k·‖₂`, read off the
    ///    row itself, so `D̃_kk` is a positively-homogeneous degree-1 function of
    ///    `H_k·` exactly as `D_kk` is. `D̃` therefore stays EXACTLY degree-one
    ///    homogeneous in `scale = λ/τ²` (and in the #991 row weight), which is
    ///    what keeps `∂B/∂ρ_sparse` on its existing seam. A fixed absolute
    ///    `sqrt(H² + ε²)` would break that homogeneity AND inject curvature into
    ///    dead atoms whose true row is ~0.
    ///
    /// `‖H_k·‖₂²` is accumulated, and never square-rooted, in the SAME
    /// diagonal-first traversal order the envelope sum uses, so the value and its
    /// θ-adjoint differentiate one floating-point expression. The off-diagonal is
    /// grouped `scale·a_k·(a_j·bracket)` — matching
    /// [`Self::row_dense_hessian`]'s `scale·a_k·(δ_kj·… + a_j·bracket)` — rather
    /// than the flat left-to-right `scale·a_k·a_j·bracket`, which differs by an
    /// ulp and made the majorized radius here and the `H` its adjoint
    /// differentiates two operators that disagreed in the last bit. That was
    /// invisible while `D` was compared at `1e-12`; it is visible the moment
    /// `D̃ ≥ D` is asserted EXACTLY, which is the form the majorization guarantee
    /// actually takes. If the row is so
    /// deeply underflowed that `Σ_l H_kl²` flushes to zero while the entries do
    /// not, `ε_k` is exactly 0 and the sum degrades gracefully to the exact hard
    /// `Σ_j|H_kj|` — the majorization guarantee is unconditional, and smoothing
    /// switches itself off only where the row's curvature is below the square
    /// root of the subnormal range and therefore invisible to `log|B|` anyway.
    pub fn psd_majorizer_abs_row_sums(&self, row: &[f64], scale: f64) -> Vec<f64> {
        let a = self.softmax_row(row);
        let k = self.k_atoms;
        let l: Vec<f64> = (0..k).map(|i| entropy_log_plus_one(a[i])).collect();
        let m: f64 = (0..k).map(|i| a[i] * l[i]).sum();
        let eps0 = Self::soft_abs_temperature(k);
        let eps0_sq = eps0 * eps0;
        let mut d = vec![0.0_f64; k];
        for kk in 0..k {
            // Diagonal entry H_kk.
            let h_kk = scale * a[kk] * ((m - l[kk] - 1.0) + a[kk] * (2.0 * l[kk] + 1.0 - 2.0 * m));
            // Pass 1: the row's own squared curvature scale ‖H_k·‖₂².
            let mut sum_sq = h_kk * h_kk;
            for jj in 0..k {
                if jj == kk {
                    continue;
                }
                let h_kj = scale * a[kk] * (a[jj] * (l[kk] + l[jj] + 1.0 - 2.0 * m));
                sum_sq += h_kj * h_kj;
            }
            // Pass 2: the soft-abs row sum at that scale, ε_k² = ε₀²·‖H_k·‖₂².
            let eps_sq = eps0_sq * sum_sq;
            let mut acc = soft_abs_squared_scale(h_kk, eps_sq);
            // Off-diagonal entries H_kj, j ≠ k.
            for jj in 0..k {
                if jj == kk {
                    continue;
                }
                let h_kj = scale * a[kk] * (a[jj] * (l[kk] + l[jj] + 1.0 - 2.0 * m));
                acc += soft_abs_squared_scale(h_kj, eps_sq);
            }
            d[kk] = acc;
        }
        d
    }

    /// Exact per-row dense softmax-entropy Hessian wrt the row's logits (#1038),
    /// scaled by `scale = λ/τ²`. Returns the symmetric `K×K` block
    ///
    /// ```text
    ///   H_kj = scale·a_k·[ δ_kj·(m − L_k − 1) + a_j·(L_k + L_j + 1 − 2m) ],
    ///   L_k = ln a_k + 1,   m = Σ_r a_r L_r,
    /// ```
    ///
    /// whose diagonal coincides with [`AnalyticPenalty::hessian_diag`] and whose
    /// quadratic form coincides with [`AnalyticPenalty::hvp`]. This is the dense
    /// block the Arrow-Schur row factor stores so the criterion's `log|H|` and
    /// the #1006 θ-adjoint differentiate the SAME operator (not just its
    /// diagonal). The entropy block alone is gauge-null (`H·𝟙 = 0`, softmax
    /// shift-invariance); callers must add it to the gauge-breaking data-fit
    /// row block before factoring — never factor it in isolation.
    #[must_use]
    pub fn row_dense_hessian(&self, row_logits: &[f64], scale: f64) -> Array2<f64> {
        let k = self.k_atoms;
        let a = self.softmax_row(row_logits);
        let l: Vec<f64> = (0..k).map(|i| entropy_log_plus_one(a[i])).collect();
        let m: f64 = (0..k).map(|i| a[i] * l[i]).sum();
        let mut h = Array2::<f64>::zeros((k, k));
        for kk in 0..k {
            for jj in 0..k {
                let indicator = if kk == jj { 1.0 } else { 0.0 };
                h[[kk, jj]] = scale
                    * a[kk]
                    * (indicator * (m - l[kk] - 1.0) + a[jj] * (l[kk] + l[jj] + 1.0 - 2.0 * m));
            }
        }
        h
    }

    /// Derivative of the exact per-row dense entropy Hessian
    /// [`Self::row_dense_hessian`] with respect to a single row logit `z_w`,
    /// scaled by `scale = λ/τ²`. Returns the symmetric `K×K` block
    /// `∂H_kj/∂z_w`, the third-derivative tensor slice the #1006 θ-adjoint
    /// contracts against the row's selected inverse. Built from the SAME
    /// `(a, L, m)` as [`Self::row_dense_hessian`] (`∂a_r/∂z_w = a_r(δ_rw − a_w)/τ`),
    /// so value, logdet and adjoint stay on one branch.
    #[must_use]
    pub fn row_dense_hessian_logit_derivative(
        &self,
        row_logits: &[f64],
        scale: f64,
        w: usize,
    ) -> Array2<f64> {
        let k = self.k_atoms;
        let inv_tau = 1.0 / self.temperature;
        let a = self.softmax_row(row_logits);
        let l: Vec<f64> = (0..k).map(|i| entropy_log_plus_one(a[i])).collect();
        let m: f64 = (0..k).map(|i| a[i] * l[i]).sum();
        // ∂a_r/∂z_w = a_r (δ_rw − a_w)/τ ; ∂L_r/∂z_w = (∂a_r/∂z_w)/a_r.
        let da: Vec<f64> = (0..k)
            .map(|r| a[r] * (if r == w { 1.0 } else { 0.0 } - a[w]) * inv_tau)
            .collect();
        let dl: Vec<f64> = (0..k)
            .map(|r| if a[r] > 0.0 { da[r] / a[r] } else { 0.0 })
            .collect();
        let dm: f64 = (0..k).map(|r| da[r] * l[r] + a[r] * dl[r]).sum();
        let mut dh = Array2::<f64>::zeros((k, k));
        for kk in 0..k {
            for jj in 0..k {
                let indicator = if kk == jj { 1.0 } else { 0.0 };
                // bracket = δ_kj(m − L_k − 1) + a_j(L_k + L_j + 1 − 2m).
                let bracket =
                    indicator * (m - l[kk] - 1.0) + a[jj] * (l[kk] + l[jj] + 1.0 - 2.0 * m);
                let dbracket = indicator * (dm - dl[kk])
                    + da[jj] * (l[kk] + l[jj] + 1.0 - 2.0 * m)
                    + a[jj] * (dl[kk] + dl[jj] - 2.0 * dm);
                dh[[kk, jj]] = scale * (da[kk] * bracket + a[kk] * dbracket);
            }
        }
        dh
    }

    /// Per-row **Gershgorin diagonal majorizer** `D̃` of the exact softmax-entropy
    /// Hessian [`Self::row_dense_hessian`], scaled by `scale = λ/τ²`. Returns the
    /// `K×K` diagonal block `diag(D̃_0, …, D̃_{K−1})` with
    /// `D̃_kk = Σ_j σ_{ε_k}(H_kj) ≥ Σ_j |H_kj|` — the smooth soft-abs envelope of
    /// the Gershgorin radius (#1419 majorizer, #2339 smoothing; the derivation
    /// and its four guarantees are on [`Self::psd_majorizer_abs_row_sums`]).
    ///
    /// Unlike the Fisher metric [`Self::row_fisher_metric`] — which is PSD but
    /// does NOT satisfy `G ⪰ H_entropy` (counterexample `a=(0.95,0.05)`,
    /// `λ=τ=1`: `G₁₁=0.0475 < H₁₁=0.0784`) — this `D̃` is a genuine Loewner
    /// majorizer. The hard radius `D_kk = Σ_j|H_kj|` is diagonally dominant over
    /// `H` (`D_kk − H_kk = |H_kk|−H_kk + Σ_{j≠k}|H_kj| ≥ Σ_{j≠k}|(D−H)_kj|`), so
    /// `D − H ⪰ 0` and `D ⪰ 0`; the envelope only ever raises each term
    /// (`σ_ε ≥ |·|`), so `D̃ − H = (D̃ − D) + (D − H)` is a sum of two PSD
    /// matrices and `D̃ ⪰ D ⪰ H`, `D̃ ⪰ D ⪰ 0`. It therefore both keeps the
    /// assembled evidence block PD (the property the entropy block needs so the
    /// Faddeev–Popov deflation never fires) AND actually majorizes the entropy
    /// curvature, which the Fisher surrogate did not. The criterion's `log|H|`,
    /// its θ-adjoint [`Self::row_psd_majorizer_logit_derivative`], and the
    /// assembled Hessian all differentiate this SAME operator `D̃`, keeping value
    /// and adjoint on one exact branch.
    #[must_use]
    pub fn row_psd_majorizer(&self, row_logits: &[f64], scale: f64) -> Array2<f64> {
        let k = self.k_atoms;
        let d = self.psd_majorizer_abs_row_sums(row_logits, scale);
        let mut out = Array2::<f64>::zeros((k, k));
        for kk in 0..k {
            out[[kk, kk]] = d[kk];
        }
        out
    }

    /// Derivative of the per-row Gershgorin majorizer [`Self::row_psd_majorizer`]
    /// with respect to a single row logit `z_w`, scaled by `scale = λ/τ²`.
    /// Returns the `K×K` diagonal block `diag(∂D̃_0/∂z_w, …)`, where `H` is the
    /// exact entropy Hessian [`Self::row_dense_hessian`] and `Ḣ_kj = ∂H_kj/∂z_w`
    /// is [`Self::row_dense_hessian_logit_derivative`]. Built from the SAME
    /// `(a, L, m)` derivative convention as the dense Hessian derivative, so the
    /// θ-adjoint differentiates the SAME `D̃` the assembly added.
    ///
    /// # Derivation (#2339)
    ///
    /// The majorized radius is `D̃_kk = Σ_j s_kj` with
    /// `s_kj = sqrt(H_kj² + ε₀²·r_k²)` and `r_k² = Σ_l H_kl²`
    /// ([`Self::psd_majorizer_abs_row_sums`]). Differentiating the square root
    /// once, and using `½·∂r_k²/∂z_w = Σ_l H_kl·Ḣ_kl`:
    ///
    /// ```text
    ///   ∂s_kj/∂z_w = [ H_kj·Ḣ_kj + ε₀²·Σ_l H_kl·Ḣ_kl ] / s_kj
    ///   ∂D̃_kk/∂z_w = Σ_j (H_kj/s_kj)·Ḣ_kj  +  ε₀²·G·Σ_j (1/s_kj),
    ///   G = Σ_l H_kl·Ḣ_kl.
    /// ```
    ///
    /// The soft sign `H_kj/s_kj ∈ (−1,1)` replaces the discontinuous `sign(H_kj)`
    /// of the hard radius (`ε₀ → 0` recovers it, term by term); the second term
    /// is the chain contribution of the row's OWN scale `r_k` and is what makes
    /// this the exact derivative of the operator actually installed — dropping it
    /// would reintroduce precisely the objective↔gradient desync the smoothing
    /// exists to remove. `s_kj` is evaluated through the same
    /// [`soft_abs_squared_scale`] seam as the value, so where that seam's
    /// rounding-direction guard binds this returns `sign(H_kj)·Ḣ_kj` — exactly
    /// the derivative of the value as implemented. `s_kj = 0` occurs only for a
    /// numerically zero row (`H_k· ≡ 0`, hence `Ḣ_k· ≡ 0`) and contributes
    /// nothing, the same exact-zero continuation the value uses.
    #[must_use]
    pub fn row_psd_majorizer_logit_derivative(
        &self,
        row_logits: &[f64],
        scale: f64,
        w: usize,
    ) -> Array2<f64> {
        let k = self.k_atoms;
        let h = self.row_dense_hessian(row_logits, scale);
        let dh = self.row_dense_hessian_logit_derivative(row_logits, scale, w);
        let eps0 = Self::soft_abs_temperature(k);
        let eps0_sq = eps0 * eps0;
        let mut out = Array2::<f64>::zeros((k, k));
        for kk in 0..k {
            // Pass 1: ‖H_k·‖₂² and G = Σ_l H_kl·Ḣ_kl, accumulated in the same
            // diagonal-first order the value uses so both differentiate one
            // floating-point expression.
            let mut sum_sq = h[[kk, kk]] * h[[kk, kk]];
            let mut cross = h[[kk, kk]] * dh[[kk, kk]];
            for jj in 0..k {
                if jj == kk {
                    continue;
                }
                sum_sq += h[[kk, jj]] * h[[kk, jj]];
                cross += h[[kk, jj]] * dh[[kk, jj]];
            }
            // Pass 2: the soft-sign contraction and the reciprocal-scale sum.
            let eps_sq = eps0_sq * sum_sq;
            let mut acc = 0.0_f64;
            let mut inv_envelope_sum = 0.0_f64;
            let s_kk = soft_abs_squared_scale(h[[kk, kk]], eps_sq);
            if s_kk != 0.0 {
                acc += (h[[kk, kk]] / s_kk) * dh[[kk, kk]];
                inv_envelope_sum += 1.0 / s_kk;
            }
            for jj in 0..k {
                if jj == kk {
                    continue;
                }
                let s_kj = soft_abs_squared_scale(h[[kk, jj]], eps_sq);
                if s_kj == 0.0 {
                    continue;
                }
                acc += (h[[kk, jj]] / s_kj) * dh[[kk, jj]];
                inv_envelope_sum += 1.0 / s_kj;
            }
            out[[kk, kk]] = acc + eps0_sq * cross * inv_envelope_sum;
        }
        out
    }

    /// Per-row softmax **Fisher-information metric** `G = scale·(diag(a) − a aᵀ)`
    /// over the row's logits, with `a = softmax(row_logits)` and
    /// `scale = λ/τ²` (#1190). Returns the symmetric `K×K` block
    ///
    /// ```text
    ///   G_kj = scale·a_k·(δ_kj − a_j).
    /// ```
    ///
    /// `G` is a covariance/Gram matrix, hence exactly PSD and smooth in the
    /// logits. It is the Fisher-information metric of the row softmax, NOT a
    /// curvature majorizer of the entropy Hessian: `G − H_entropy` can be
    /// indefinite (#1419: `K=2`, `a=(0.95,0.05)`, `λ=τ=1` gives `G₁₁=0.0475 <
    /// H₁₁=0.0784`, so `G ⋡ H`). The genuine Loewner majorizer the assembled
    /// evidence block now uses is [`Self::row_psd_majorizer`]
    /// (`D̃_kk = Σ_j σ_ε(H_kj) ≥ Σ_j|H_kj|`, which DOES satisfy `D̃ ⪰ H` and
    /// `D̃ ⪰ 0`); this Fisher metric is retained only as a smooth PSD
    /// conditioning reference and its derivative
    /// [`Self::row_fisher_metric_logit_derivative`], and must not be presented or
    /// used as a curvature majorizer.
    #[must_use]
    pub fn row_fisher_metric(&self, row_logits: &[f64], scale: f64) -> Array2<f64> {
        let k = self.k_atoms;
        let a = self.softmax_row(row_logits);
        let mut g = Array2::<f64>::zeros((k, k));
        for kk in 0..k {
            for jj in 0..k {
                let indicator = if kk == jj { 1.0 } else { 0.0 };
                g[[kk, jj]] = scale * a[kk] * (indicator - a[jj]);
            }
        }
        g
    }

    /// Derivative of the per-row softmax Fisher metric
    /// [`Self::row_fisher_metric`] with respect to a single row logit `z_w`,
    /// scaled by `scale = λ/τ²` (#1190). Returns the symmetric `K×K` block
    /// `∂G_kj/∂z_w`, the third-derivative tensor slice the θ-adjoint contracts
    /// against the row's selected inverse so the adjoint differentiates the SAME
    /// PSD `G = scale·(diag(a) − a aᵀ)` the assembly added (value/adjoint on one
    /// branch, no deflation needed). Built from the SAME softmax derivative
    /// convention as [`Self::row_dense_hessian_logit_derivative`]
    /// (`∂a_r/∂z_w = a_r(δ_rw − a_w)/τ`). For `G_kj = scale·a_k(δ_kj − a_j)`,
    /// the product rule gives
    /// `∂G_kj/∂z_w = scale·[ (∂a_k/∂z_w)(δ_kj − a_j) − a_k(∂a_j/∂z_w) ]`.
    #[must_use]
    pub fn row_fisher_metric_logit_derivative(
        &self,
        row_logits: &[f64],
        scale: f64,
        w: usize,
    ) -> Array2<f64> {
        let k = self.k_atoms;
        let inv_tau = 1.0 / self.temperature;
        let a = self.softmax_row(row_logits);
        // ∂a_r/∂z_w = a_r (δ_rw − a_w)/τ — identical convention to the entropy
        // Hessian derivative above.
        let da: Vec<f64> = (0..k)
            .map(|r| a[r] * (if r == w { 1.0 } else { 0.0 } - a[w]) * inv_tau)
            .collect();
        let mut dg = Array2::<f64>::zeros((k, k));
        for kk in 0..k {
            for jj in 0..k {
                let indicator = if kk == jj { 1.0 } else { 0.0 };
                dg[[kk, jj]] = scale * (da[kk] * (indicator - a[jj]) - a[kk] * da[jj]);
            }
        }
        dg
    }
}

impl AnalyticPenalty for SoftmaxAssignmentSparsityPenalty {
    fn tier(&self) -> PenaltyTier {
        PenaltyTier::Psi
    }

    fn validate_rho(&self, rho: ArrayView1<'_, f64>) -> Result<(), String> {
        if rho.len() != 1 {
            return Err(format!(
                "softmax assignment sparsity rho length {} != 1",
                rho.len()
            ));
        }
        resolve_learnable_weight(self.weight, rho[0])?;
        Ok(())
    }

    fn rho_coordinate_domains(&self) -> Result<Vec<(f64, f64)>, String> {
        Ok(vec![
            learnable_weight_coordinate_domain(self.weight)?
                .ok_or_else(|| "softmax assignment sparsity has zero base weight".to_string())?,
        ])
    }

    fn value(&self, target: ArrayView1<'_, f64>, rho: ArrayView1<'_, f64>) -> f64 {
        let lambda = validated_learnable_weight(self.weight, rho[0]);
        let n = target.len() / self.k_atoms;
        let values: Vec<f64> = target.iter().copied().collect();
        let mut acc = 0.0;
        for row in 0..n {
            let start = row * self.k_atoms;
            let a = self.softmax_row(&values[start..start + self.k_atoms]);
            let w_row = self.row_weight(row);
            for v in a {
                if v > 0.0 {
                    acc += -w_row * v * v.ln();
                }
            }
        }
        lambda * acc
    }

    fn grad_target(&self, target: ArrayView1<'_, f64>, rho: ArrayView1<'_, f64>) -> Array1<f64> {
        let lambda = validated_learnable_weight(self.weight, rho[0]);
        let n = target.len() / self.k_atoms;
        let values: Vec<f64> = target.iter().copied().collect();
        let mut out = Array1::<f64>::zeros(target.len());
        let inv_tau = 1.0 / self.temperature;
        for row in 0..n {
            let start = row * self.k_atoms;
            let a = self.softmax_row(&values[start..start + self.k_atoms]);
            let w_row = self.row_weight(row);
            let mut d_h_da = vec![0.0; self.k_atoms];
            let mut mean = 0.0;
            for k in 0..self.k_atoms {
                d_h_da[k] = -lambda * entropy_log_plus_one(a[k]);
                mean += a[k] * d_h_da[k];
            }
            for k in 0..self.k_atoms {
                out[start + k] = w_row * a[k] * (d_h_da[k] - mean) * inv_tau;
            }
        }
        out
    }

    fn hessian_diag(
        &self,
        target: ArrayView1<'_, f64>,
        rho: ArrayView1<'_, f64>,
    ) -> Option<Array1<f64>> {
        assert_eq!(rho.len(), 1, "softmax entropy expects one rho parameter");
        assert!(
            rho.iter().all(|value| value.is_finite()),
            "softmax entropy rho must be finite"
        );
        assert_eq!(
            target.len() % self.k_atoms,
            0,
            "softmax entropy target length must be divisible by k_atoms"
        );
        // Closed-form diagonal of the softmax-entropy Hessian wrt logits.
        // Derived by probing the row-dense HVP with the unit vector e_k:
        // for a row with softmax weights a_k and L_k = ln a_k + 1,
        //   H_kk = (lambda / tau^2) * a_k *
        //          ((1 - 2 a_k) * (E_a[L] - L_k) + a_k - 1).
        // This matches `hvp(...) . e_k` analytically (see derivation in the
        // bug-fix comment on `hvp`) and gives Newton/Arrow-Schur callers a
        // principled diagonal surrogate without per-row dense factorization.
        let lambda = validated_learnable_weight(self.weight, rho[0]);
        let inv_tau = 1.0 / self.temperature;
        let scale = lambda * inv_tau * inv_tau;
        let n = target.len() / self.k_atoms;
        let values: Vec<f64> = target.iter().copied().collect();
        let mut out = Array1::<f64>::zeros(target.len());
        for row in 0..n {
            let start = row * self.k_atoms;
            let a = self.softmax_row(&values[start..start + self.k_atoms]);
            let w_row = self.row_weight(row);
            let mut mean_log_plus_one = 0.0;
            for k in 0..self.k_atoms {
                mean_log_plus_one += a[k] * entropy_log_plus_one(a[k]);
            }
            for k in 0..self.k_atoms {
                let log_plus_one = entropy_log_plus_one(a[k]);
                let term = (1.0 - 2.0 * a[k]) * (mean_log_plus_one - log_plus_one) + a[k] - 1.0;
                out[start + k] = w_row * scale * a[k] * term;
            }
        }
        Some(out)
    }

    fn hvp(
        &self,
        target: ArrayView1<'_, f64>,
        rho: ArrayView1<'_, f64>,
        v: ArrayView1<'_, f64>,
    ) -> Array1<f64> {
        /*
        Softmax entropy is not coordinate-separable in logits. The old
        `hessian_diag` returned λ p_k(1-p_k)/τ², which is only the softmax
        Jacobian diagonal and omits the entropy curvature and all cross-logit
        terms. For H(p(z)), p'=p*(v-E_p[v])/τ and
        (log p_k + 1)'=(v_k-E_p[v])/τ. Differentiating
        g_k=λ p_k(E_p[log p + 1]-(log p_k+1))/τ gives the row-dense product
        below. `hessian_diag` returns the analytic diagonal extracted from
        this HVP by setting v = e_k row-by-row.
        */
        let lambda = validated_learnable_weight(self.weight, rho[0]);
        assert_eq!(target.len(), v.len(), "hvp dimension mismatch");
        let n = target.len() / self.k_atoms;
        let values: Vec<f64> = target.iter().copied().collect();
        let mut out = Array1::<f64>::zeros(target.len());
        let inv_tau = 1.0 / self.temperature;
        let scale = lambda * inv_tau * inv_tau;
        for row in 0..n {
            let start = row * self.k_atoms;
            let a = self.softmax_row(&values[start..start + self.k_atoms]);
            let w_row = self.row_weight(row);
            let mut mean_log_plus_one = 0.0;
            let mut mean_v = 0.0;
            for k in 0..self.k_atoms {
                mean_log_plus_one += a[k] * entropy_log_plus_one(a[k]);
                mean_v += a[k] * v[start + k];
            }
            let mut mean_centered_v_log_plus_one = 0.0;
            for k in 0..self.k_atoms {
                let centered_v = v[start + k] - mean_v;
                mean_centered_v_log_plus_one += a[k] * centered_v * entropy_log_plus_one(a[k]);
            }
            for k in 0..self.k_atoms {
                let log_plus_one = entropy_log_plus_one(a[k]);
                let centered_v = v[start + k] - mean_v;
                out[start + k] = w_row
                    * scale
                    * a[k]
                    * (centered_v * (mean_log_plus_one - log_plus_one - 1.0)
                        + mean_centered_v_log_plus_one);
            }
        }
        out
    }

    fn psd_majorizer_diag(
        &self,
        target: ArrayView1<'_, f64>,
        rho: ArrayView1<'_, f64>,
    ) -> Option<Array1<f64>> {
        assert_eq!(rho.len(), 1, "softmax entropy expects one rho parameter");
        assert_eq!(
            target.len() % self.k_atoms,
            0,
            "softmax entropy target length must be divisible by k_atoms"
        );
        // Entropy minimization is nonconvex: the exact per-row Hessian is dense
        // and indefinite, so the convex-only trait default (which returns the
        // raw indefinite `hessian_diag`) violates the `B ⪰ 0` contract and is a
        // diagonal masquerading as a dense operator. Replace it with the
        // Gershgorin / diagonal-dominance majorizer of the dense per-row block
        // (see `psd_majorizer_abs_row_sums`): a genuine PSD diagonal with
        // `D ⪰ H` and `D ⪰ 0`. Coordinate-indexed, so the inherited
        // `psd_majorizer_hvp` applies `D` as a diagonal operator consistently.
        let lambda = validated_learnable_weight(self.weight, rho[0]);
        let inv_tau = 1.0 / self.temperature;
        let scale = lambda * inv_tau * inv_tau;
        let n = target.len() / self.k_atoms;
        let values: Vec<f64> = target.iter().copied().collect();
        let mut out = Array1::<f64>::zeros(target.len());
        for row in 0..n {
            let start = row * self.k_atoms;
            let w_row = self.row_weight(row);
            let d = self.psd_majorizer_abs_row_sums(&values[start..start + self.k_atoms], scale);
            for k in 0..self.k_atoms {
                out[start + k] = w_row * d[k];
            }
        }
        Some(out)
    }

    fn grad_rho(&self, target: ArrayView1<'_, f64>, rho: ArrayView1<'_, f64>) -> Array1<f64> {
        Array1::from_vec(vec![self.value(target, rho)])
    }

    fn rho_count(&self) -> usize {
        1
    }

    fn name(&self) -> &str {
        "softmax_assignment_sparsity"
    }

    impl_scalar_apply_schedule!(weight);
}

impl SparsityPenalty {
    #[must_use = "build error must be handled"]
    pub fn smoothed_l1(target_tier: PenaltyTier, eps: f64) -> Result<Self, String> {
        if !(eps.is_finite() && eps > 0.0) {
            return Err(format!(
                "SparsityPenalty::smoothed_l1 requires eps > 0 \
                 (Hessian / gradient have a `1/sqrt(x² + eps²)` factor that needs eps > 0 \
                 for differentiability at x = 0); got eps = {eps}"
            ));
        }
        Ok(Self {
            target_tier,
            kind: SparsityKind::SmoothedL1 { eps },
            weight: 1.0,
            weight_schedule: None,
            learnable_smoothing: false,
        })
    }

    #[must_use = "build error must be handled"]
    pub fn log(target_tier: PenaltyTier, delta: f64) -> Result<Self, String> {
        if !(delta.is_finite() && delta > 0.0) {
            return Err(format!(
                "SparsityPenalty::log requires delta > 0 \
                 (the log-sparsifier is log(1 + x²/δ²), undefined at δ = 0); \
                 got delta = {delta}"
            ));
        }
        Ok(Self {
            target_tier,
            kind: SparsityKind::Log { delta },
            weight: 1.0,
            weight_schedule: None,
            learnable_smoothing: false,
        })
    }

    /// Hoyer scale-invariant sparsifier. Requires a target of length > 1
    /// because the normalized form divides by `sqrt(n) - 1`.
    #[must_use]
    pub fn hoyer(target_tier: PenaltyTier) -> Self {
        Self {
            target_tier,
            kind: SparsityKind::Hoyer,
            weight: 1.0,
            weight_schedule: None,
            learnable_smoothing: false,
        }
    }

    impl_with_weight_schedule!(weight);

    #[must_use = "invalid learnable-smoothing requests must be handled"]
    pub fn with_learnable_smoothing(mut self) -> Result<Self, String> {
        if matches!(self.kind, SparsityKind::Hoyer) {
            return Err("Hoyer sparsity has no smoothing coordinate to learn".to_string());
        }
        // Coordinate 0 is the strength and coordinate 1 is the optional
        // smoothing log-scale. Do not accept an arbitrary index: rho_count is
        // exactly two in this state, so any other index is structurally
        // impossible and would defer a builder error into evaluator indexing.
        self.learnable_smoothing = true;
        Ok(self)
    }

    #[must_use]
    pub fn learns_smoothing(&self) -> bool {
        self.learnable_smoothing
    }

    /// Resolve `(strength, eps_or_delta)` from the current ρ view.
    fn resolved(&self, rho: ArrayView1<'_, f64>) -> (f64, f64) {
        let strength = validated_learnable_weight(self.weight, rho[0]);
        let smoothing = match (self.learnable_smoothing, self.kind) {
            // The owning seam validates this log-smoothing coordinate before
            // exact exponentiation, so it stays positive without a saturated
            // tail or value/derivative mismatch.
            (true, _) => validated_exp_log_strength(rho[1]),
            (false, SparsityKind::SmoothedL1 { eps }) => eps,
            (false, SparsityKind::Log { delta }) => delta,
            (false, SparsityKind::Hoyer) => 0.0,
        };
        (strength, smoothing)
    }
}

impl AnalyticPenalty for SparsityPenalty {
    fn tier(&self) -> PenaltyTier {
        self.target_tier
    }

    fn validate_rho(&self, rho: ArrayView1<'_, f64>) -> Result<(), String> {
        if rho.len() != self.rho_count() {
            return Err(format!(
                "sparsity rho length {} != declared {}",
                rho.len(),
                self.rho_count()
            ));
        }
        resolve_learnable_weight(self.weight, rho[0])?;
        if self.learnable_smoothing {
            checked_exp_log_strength(rho[1]).map_err(|error| error.to_string())?;
        }
        Ok(())
    }

    fn rho_coordinate_domains(&self) -> Result<Vec<(f64, f64)>, String> {
        let mut domains = vec![(LOG_STRENGTH_MIN, LOG_STRENGTH_MAX); self.rho_count()];
        domains[0] = learnable_weight_coordinate_domain(self.weight)?
            .ok_or_else(|| "sparsity has zero base weight".to_string())?;
        Ok(domains)
    }

    fn value(&self, target: ArrayView1<'_, f64>, rho: ArrayView1<'_, f64>) -> f64 {
        let (lam, smooth) = self.resolved(rho);
        match self.kind {
            SparsityKind::SmoothedL1 { .. } => {
                let mut acc = 0.0;
                for &x in target.iter() {
                    acc += (x * x + smooth * smooth).sqrt();
                }
                lam * acc
            }
            SparsityKind::Hoyer => {
                // Normalized anti-sparsity penalty
                //   P(x) = (||x||_1 / ||x||_2 - 1) / (sqrt(n) - 1)
                // maps [1, sqrt(n)] -> [0, 1]. A perfectly dense
                // equal-magnitude vector hits ||x||_1/||x||_2 = sqrt(n),
                // so P = 1; a 1-sparse vector has ratio 1, so P = 0
                // (sparse vectors minimize the penalty).
                let n = target.len() as f64;
                assert!(n > 1.0, "Hoyer requires n > 1");
                let l1: f64 = target.iter().map(|x| x.abs()).sum();
                let l2: f64 = target.iter().map(|x| x * x).sum::<f64>().sqrt();
                if l2 == 0.0 {
                    return 0.0;
                }
                let h = (l1 / l2 - 1.0) / (n.sqrt() - 1.0);
                lam * h
            }
            SparsityKind::Log { .. } => {
                let mut acc = 0.0;
                let d2 = smooth * smooth;
                for &x in target.iter() {
                    acc += (1.0 + x * x / d2).ln();
                }
                lam * acc
            }
        }
    }

    fn grad_target(&self, target: ArrayView1<'_, f64>, rho: ArrayView1<'_, f64>) -> Array1<f64> {
        let (lam, smooth) = self.resolved(rho);
        let mut g = Array1::<f64>::zeros(target.len());
        match self.kind {
            SparsityKind::SmoothedL1 { .. } => {
                let eps2 = smooth * smooth;
                for (i, &x) in target.iter().enumerate() {
                    g[i] = lam * x / (x * x + eps2).sqrt();
                }
            }
            SparsityKind::Hoyer => {
                // P(x) = A · (L1/L2 - 1), A = lam / (sqrt(n) - 1).
                // ∂P/∂x_i = A · (sign(x_i)/L2 - L1 · x_i / L2³).
                let n = target.len() as f64;
                assert!(n > 1.0, "Hoyer requires n > 1");
                let l1: f64 = target.iter().map(|x| x.abs()).sum();
                let l2: f64 = target.iter().map(|x| x * x).sum::<f64>().sqrt();
                if l2 == 0.0 {
                    return g;
                }
                let denom = n.sqrt() - 1.0;
                let a = lam / denom;
                let inv_l2 = 1.0 / l2;
                let inv_l2_cubed = inv_l2 * inv_l2 * inv_l2;
                for (i, &x) in target.iter().enumerate() {
                    let sgn = if x > 0.0 {
                        1.0
                    } else if x < 0.0 {
                        -1.0
                    } else {
                        0.0
                    };
                    g[i] = a * (sgn * inv_l2 - l1 * x * inv_l2_cubed);
                }
            }
            SparsityKind::Log { .. } => {
                let d2 = smooth * smooth;
                for (i, &x) in target.iter().enumerate() {
                    g[i] = lam * 2.0 * x / (d2 + x * x);
                }
            }
        }
        g
    }

    fn hessian_diag(
        &self,
        target: ArrayView1<'_, f64>,
        rho: ArrayView1<'_, f64>,
    ) -> Option<Array1<f64>> {
        let (lam, smooth) = self.resolved(rho);
        match self.kind {
            SparsityKind::SmoothedL1 { .. } => {
                let mut d = Array1::<f64>::zeros(target.len());
                let eps2 = smooth * smooth;
                for (i, &x) in target.iter().enumerate() {
                    let r = (x * x + eps2).sqrt();
                    d[i] = lam * eps2 / (r * r * r);
                }
                Some(d)
            }
            SparsityKind::Log { .. } => {
                let mut d = Array1::<f64>::zeros(target.len());
                // The EXACT second derivative of λ log(1 + x²/δ²):
                //   d/dx [ 2λx/(δ²+x²) ] = 2λ(δ² − x²)/(δ² + x²)²,
                // which is NEGATIVE for |x| > δ — Log is nonconvex. This is
                // the genuine Hessian diagonal and exactly differentiates
                // `grad_target`. PSD consumers (Newton block, preconditioner,
                // `log_det_plus_λI`, FrozenAnalyticPenaltyOp) must instead
                // route through `psd_majorizer_diag`/`psd_majorizer_hvp`,
                // which expose the IRLS/MM surrogate `2λ/(δ²+x²)`.
                let d2 = smooth * smooth;
                for (i, &x) in target.iter().enumerate() {
                    let denom = d2 + x * x;
                    d[i] = lam * 2.0 * (d2 - x * x) / (denom * denom);
                }
                Some(d)
            }
            // Hoyer's Hessian is DENSE and NOT generally PSD (Hoyer is a
            // nonconvex sparsifier). We cannot return a meaningful diagonal
            // that would be safe to use as a preconditioner / Newton block
            // through the standard `hessian_diag` path, so we return `None`
            // and force callers through `hvp`. See `hvp` below for the exact
            // dense-Hessian-vector product.
            SparsityKind::Hoyer => None,
        }
    }

    fn hvp(
        &self,
        target: ArrayView1<'_, f64>,
        rho: ArrayView1<'_, f64>,
        v: ArrayView1<'_, f64>,
    ) -> Array1<f64> {
        // For SmoothedL1/Log/Hoyer we route through the closed-form Hessian.
        // SmoothedL1 and Log have purely diagonal Hessians and would
        // ordinarily reach the diagonal branch of the default `hvp`; we
        // override here to also serve Hoyer (whose Hessian is dense
        // rank-1-plus-diagonal).
        let (lam, smooth) = self.resolved(rho);
        let n_target = target.len();
        assert_eq!(v.len(), n_target, "hvp dimension mismatch");
        match self.kind {
            SparsityKind::SmoothedL1 { .. } => {
                let mut out = Array1::<f64>::zeros(n_target);
                let eps2 = smooth * smooth;
                for (i, &x) in target.iter().enumerate() {
                    let r = (x * x + eps2).sqrt();
                    out[i] = lam * eps2 / (r * r * r) * v[i];
                }
                out
            }
            SparsityKind::Log { .. } => {
                // EXACT Hessian-vector product: the Log Hessian is diagonal
                // with entries 2λ(δ²−x²)/(δ²+x²)², so (Hv)_i = h_i v_i. This
                // is the genuine second derivative (indefinite for |x|>δ).
                // PSD consumers use `psd_majorizer_hvp` for the IRLS/MM
                // surrogate 2λ/(δ²+x²) instead.
                let mut out = Array1::<f64>::zeros(n_target);
                let d2 = smooth * smooth;
                for (i, &x) in target.iter().enumerate() {
                    let denom = d2 + x * x;
                    out[i] = lam * 2.0 * (d2 - x * x) / (denom * denom) * v[i];
                }
                out
            }
            SparsityKind::Hoyer => {
                // P(x) = A · (L1/L2 - 1), A = lam / (sqrt(n) - 1).
                // H_ij = A · [ -s_i x_j/L2³ - x_i s_j/L2³
                //              - L1 δ_ij/L2³ + 3 L1 x_i x_j/L2⁵ ]
                // (Hv)_i = A · [ -s_i (xᵀv)/L2³ - x_i (sᵀv)/L2³
                //                - L1 v_i/L2³ + 3 L1 x_i (xᵀv)/L2⁵ ]
                let n = n_target as f64;
                assert!(n > 1.0, "Hoyer requires n > 1");
                let l1: f64 = target.iter().map(|x| x.abs()).sum();
                let l2: f64 = target.iter().map(|x| x * x).sum::<f64>().sqrt();
                let mut out = Array1::<f64>::zeros(n_target);
                if l2 == 0.0 {
                    return out;
                }
                let a = lam / (n.sqrt() - 1.0);
                let inv_l2_cubed = 1.0 / (l2 * l2 * l2);
                let inv_l2_5 = inv_l2_cubed / (l2 * l2);
                let mut x_dot_v = 0.0;
                let mut s_dot_v = 0.0;
                for i in 0..n_target {
                    let xi = target[i];
                    let si = if xi > 0.0 {
                        1.0
                    } else if xi < 0.0 {
                        -1.0
                    } else {
                        0.0
                    };
                    x_dot_v += xi * v[i];
                    s_dot_v += si * v[i];
                }
                for i in 0..n_target {
                    let xi = target[i];
                    let si = if xi > 0.0 {
                        1.0
                    } else if xi < 0.0 {
                        -1.0
                    } else {
                        0.0
                    };
                    out[i] = a
                        * (-si * x_dot_v * inv_l2_cubed
                            - xi * s_dot_v * inv_l2_cubed
                            - l1 * v[i] * inv_l2_cubed
                            + 3.0 * l1 * xi * x_dot_v * inv_l2_5);
                }
                out
            }
        }
    }

    fn psd_majorizer_diag(
        &self,
        target: ArrayView1<'_, f64>,
        rho: ArrayView1<'_, f64>,
    ) -> Option<Array1<f64>> {
        let (lam, smooth) = self.resolved(rho);
        match self.kind {
            // SmoothedL1 is convex: the majorizer equals the exact Hessian.
            SparsityKind::SmoothedL1 { .. } => self.hessian_diag(target, rho),
            // Log is nonconvex; expose the IRLS/MM re-weighted-ℓ₂ surrogate
            //   2λ/(δ²+x²) ⪰ 2λ(δ²−x²)/(δ²+x²)²,
            // strictly positive, agreeing with the exact Hessian at x = 0.
            SparsityKind::Log { .. } => {
                let mut d = Array1::<f64>::zeros(target.len());
                let d2 = smooth * smooth;
                for (i, &x) in target.iter().enumerate() {
                    d[i] = lam * 2.0 / (d2 + x * x);
                }
                Some(d)
            }
            // Hoyer's Hessian is dense; no diagonal majorizer. Callers fall
            // back to the exact dense `hvp` through `psd_majorizer_hvp`.
            SparsityKind::Hoyer => None,
        }
    }

    fn grad_rho(&self, target: ArrayView1<'_, f64>, rho: ArrayView1<'_, f64>) -> Array1<f64> {
        // Strength axis: ∂P/∂ρ_strength = P (chain rule through exp).
        // ε axis (if owned): ∂P/∂ρ_eps = ε · ∂P/∂ε.
        let n_rho = self.rho_count();
        let mut out = Array1::<f64>::zeros(n_rho);
        let p_val = self.value(target, rho);
        out[0] = p_val;
        if self.learnable_smoothing {
            let (lam, smooth) = self.resolved(rho);
            let mut dp_deps = 0.0;
            match self.kind {
                SparsityKind::SmoothedL1 { .. } => {
                    for &x in target.iter() {
                        dp_deps += smooth / (x * x + smooth * smooth).sqrt();
                    }
                    dp_deps *= lam;
                }
                SparsityKind::Log { .. } => {
                    // d/dδ log(1 + x²/δ²) = -2 x² / (δ (δ² + x²))
                    let d2 = smooth * smooth;
                    for &x in target.iter() {
                        dp_deps += -2.0 * x * x / (smooth * (d2 + x * x));
                    }
                    dp_deps *= lam;
                }
                SparsityKind::Hoyer => {}
            }
            // Chain through ρ_eps = log(ε)  ⇒  ∂ε/∂ρ_eps = ε.
            out[1] = smooth * dp_deps;
        }
        out
    }

    fn rho_count(&self) -> usize {
        1 + usize::from(self.learnable_smoothing)
    }

    fn name(&self) -> &str {
        "sparsity"
    }

    impl_scalar_apply_schedule!(weight);
}

// ---------------------------------------------------------------------------
// TopK activation penalty
// ---------------------------------------------------------------------------

#[derive(Debug, Clone)]
pub struct TopKActivationPenalty {
    pub target: PsiSlice,
    pub k: usize,
    pub latent_dim: usize,
    pub weight: f64,
    pub weight_schedule: Option<ScalarWeightSchedule>,
}

impl TopKActivationPenalty {
    #[must_use = "build error must be handled"]
    pub fn new(target: PsiSlice, k: usize, weight: f64) -> Result<Self, String> {
        let latent_dim = target
            .latent_dim
            .ok_or_else(|| "TopKActivationPenalty::new requires target.latent_dim".to_string())?;
        if latent_dim == 0 {
            return Err("TopKActivationPenalty::new requires latent_dim > 0".to_string());
        }
        if k == 0 || k > latent_dim {
            return Err(format!(
                "TopKActivationPenalty::new requires 0 < k <= latent_dim; got k={k}, latent_dim={latent_dim}"
            ));
        }
        if !(weight.is_finite() && weight > 0.0) {
            return Err(format!(
                "TopKActivationPenalty::new requires finite weight > 0, got {weight}"
            ));
        }
        Ok(Self {
            target,
            k,
            latent_dim,
            weight,
            weight_schedule: None,
        })
    }

    impl_with_weight_schedule!(weight);

    fn topk_mask_row(&self, target: ArrayView1<'_, f64>, row: usize, mask: &mut [bool]) {
        mask.fill(false);
        let d = self.latent_dim;
        let base = row * d;
        let mut order = (0..d).collect::<Vec<_>>();
        order.sort_by(|&a, &b| {
            target[base + b]
                .abs()
                .total_cmp(&target[base + a].abs())
                .then_with(|| a.cmp(&b))
        });
        for &axis in order.iter().take(self.k) {
            mask[axis] = true;
        }
    }
}

impl AnalyticPenalty for TopKActivationPenalty {
    fn tier(&self) -> PenaltyTier {
        PenaltyTier::Psi
    }

    fn value(&self, target: ArrayView1<'_, f64>, rho: ArrayView1<'_, f64>) -> f64 {
        assert_eq!(rho.len(), 0, "TopKActivationPenalty has no rho parameters");
        let d = self.latent_dim;
        let n_obs = target.len() / d;
        let mut mask = vec![false; d];
        let mut acc = 0.0;
        for row in 0..n_obs {
            self.topk_mask_row(target, row, &mut mask);
            let base = row * d;
            for axis in 0..d {
                if mask[axis] {
                    let v = target[base + axis];
                    acc += 0.5 * self.weight * v * v;
                }
            }
        }
        acc
    }

    fn grad_target(&self, target: ArrayView1<'_, f64>, rho: ArrayView1<'_, f64>) -> Array1<f64> {
        assert_eq!(rho.len(), 0, "TopKActivationPenalty has no rho parameters");
        let d = self.latent_dim;
        let n_obs = target.len() / d;
        let mut mask = vec![false; d];
        let mut grad = Array1::<f64>::zeros(target.len());
        for row in 0..n_obs {
            self.topk_mask_row(target, row, &mut mask);
            let base = row * d;
            for axis in 0..d {
                if mask[axis] {
                    grad[base + axis] = self.weight * target[base + axis];
                }
            }
        }
        grad
    }

    fn hessian_diag(
        &self,
        target: ArrayView1<'_, f64>,
        rho: ArrayView1<'_, f64>,
    ) -> Option<Array1<f64>> {
        assert_eq!(rho.len(), 0, "TopKActivationPenalty has no rho parameters");
        let d = self.latent_dim;
        let n_obs = target.len() / d;
        let mut mask = vec![false; d];
        let mut diag = Array1::<f64>::zeros(target.len());
        for row in 0..n_obs {
            self.topk_mask_row(target, row, &mut mask);
            let base = row * d;
            for axis in 0..d {
                if mask[axis] {
                    diag[base + axis] = self.weight;
                }
            }
        }
        Some(diag)
    }

    fn grad_rho(&self, target: ArrayView1<'_, f64>, rho: ArrayView1<'_, f64>) -> Array1<f64> {
        assert_eq!(rho.len(), 0, "TopKActivationPenalty has no rho parameters");
        assert_eq!(
            target.len() % self.latent_dim,
            0,
            "TopKActivationPenalty target length must be a multiple of latent_dim"
        );
        Array1::<f64>::zeros(0)
    }

    fn rho_count(&self) -> usize {
        0
    }

    fn name(&self) -> &str {
        "topk_activation"
    }

    impl_scalar_apply_schedule!(weight);
}

// ---------------------------------------------------------------------------
// Smooth threshold penalty
// ---------------------------------------------------------------------------

#[derive(Debug, Clone)]
pub struct SmoothThresholdPenalty {
    pub target: PsiSlice,
    pub latent_dim: usize,
    pub thresholds: Array1<f64>,
    pub weight: f64,
    pub smoothing_eps: f64,
    pub weight_schedule: Option<ScalarWeightSchedule>,
}

impl SmoothThresholdPenalty {
    #[must_use = "build error must be handled"]
    pub fn new(
        target: PsiSlice,
        thresholds: Array1<f64>,
        weight: f64,
        smoothing_eps: f64,
    ) -> Result<Self, String> {
        let latent_dim = target
            .latent_dim
            .ok_or_else(|| "SmoothThresholdPenalty::new requires target.latent_dim".to_string())?;
        if latent_dim == 0 {
            return Err("SmoothThresholdPenalty::new requires latent_dim > 0".to_string());
        }
        if thresholds.len() != latent_dim {
            return Err(format!(
                "SmoothThresholdPenalty::new thresholds length {} does not match latent_dim {latent_dim}",
                thresholds.len()
            ));
        }
        for (idx, &tau) in thresholds.iter().enumerate() {
            if !(tau.is_finite() && tau > 0.0) {
                return Err(format!(
                    "SmoothThresholdPenalty::new thresholds[{idx}] must be finite and > 0, got {tau}"
                ));
            }
        }
        if !(weight.is_finite() && weight > 0.0) {
            return Err(format!(
                "SmoothThresholdPenalty::new requires finite weight > 0, got {weight}"
            ));
        }
        if !(smoothing_eps.is_finite() && smoothing_eps > 0.0) {
            return Err(format!(
                "SmoothThresholdPenalty::new requires finite smoothing_eps > 0, got {smoothing_eps}"
            ));
        }
        Ok(Self {
            target,
            latent_dim,
            thresholds,
            weight,
            smoothing_eps,
            weight_schedule: None,
        })
    }

    impl_with_weight_schedule!(weight);

    fn threshold(&self, axis: usize, rho: ArrayView1<'_, f64>) -> f64 {
        // Resolve the exact multiplicative threshold after the owning seam has
        // validated its effective log-strength domain.
        validated_learnable_weight(self.thresholds[axis], rho[axis])
    }

    pub(crate) fn sigmoid_gate(&self, x: f64) -> f64 {
        if x >= 0.0 {
            1.0 / (1.0 + (-x).exp())
        } else {
            let ex = x.exp();
            ex / (1.0 + ex)
        }
    }

    fn true_hessian_diag_entry(&self, tau: f64, gate: f64) -> f64 {
        self.weight * tau * gate * (1.0 - gate) * (1.0 - 2.0 * gate)
            / (self.smoothing_eps * self.smoothing_eps)
    }

    fn psd_hessian_diag_entry(&self, tau: f64, gate: f64) -> f64 {
        // Genuine PSD majorizer of the indefinite exact diagonal Hessian
        //   h(g) = λτ·g(1−g)(1−2g)/ε².
        // The bare re-weighted-ℓ₂ surrogate λτ·[g(1−g)]²/ε² is ≥ 0 but only
        // dominates h in the concave region g > ½. For g < (3−√5)/2 ≈ 0.382 the
        // exact curvature is positive and strictly larger, so the square alone
        // is NOT an upper bound — the `B ⪰ ∂²P` contract is violated for exactly
        // the comfortably-below-threshold coordinates this penalty is
        // meant to suppress, costing the MM step its monotone-decrease guarantee.
        //
        // Take the elementwise max of that surrogate and the absolute exact
        // Hessian |h| = λτ·g(1−g)|1−2g|/ε². Since |h| ≥ h everywhere and ≥ 0, the
        // max is a true PSD upper bound; it equals |h| in the wings (tight where
        // the bare square failed) and keeps the surrogate's strictly-positive
        // floor near the inflection g ≈ ½ (where h ≈ 0) so the curvature block
        // never collapses to zero.
        let slope = gate * (1.0 - gate);
        let reweighted_l2 = slope * slope;
        let abs_exact = slope * (1.0 - 2.0 * gate).abs();
        self.weight * tau * reweighted_l2.max(abs_exact) / (self.smoothing_eps * self.smoothing_eps)
    }
}

/// Smooth threshold activation `φ(z) = z · σ((z − τ)/ε)` and its exact
/// derivatives:
///
///   g       = σ((z − τ)/ε)
///   φ        = z · g
///   ∂φ/∂z   = g + z · g (1 − g) / ε
///   ∂φ/∂τ   = − z · g (1 − g) / ε
#[must_use]
pub fn smooth_threshold_gate_value_grad(z: f64, tau: f64, smoothing_eps: f64) -> (f64, f64, f64) {
    let g = gam_linalg::utils::stable_logistic((z - tau) / smoothing_eps);
    let value = z * g;
    let slope = z * g * (1.0 - g) / smoothing_eps;
    let dphi_dz = g + slope;
    let dphi_dtau = -slope;
    (value, dphi_dz, dphi_dtau)
}

impl AnalyticPenalty for SmoothThresholdPenalty {
    fn tier(&self) -> PenaltyTier {
        PenaltyTier::Psi
    }

    fn validate_rho(&self, rho: ArrayView1<'_, f64>) -> Result<(), String> {
        if rho.len() != self.latent_dim {
            return Err(format!(
                "smooth-threshold rho length {} != latent dimension {}",
                rho.len(),
                self.latent_dim
            ));
        }
        for axis in 0..self.latent_dim {
            resolve_learnable_weight(self.thresholds[axis], rho[axis])?;
        }
        Ok(())
    }

    fn rho_coordinate_domains(&self) -> Result<Vec<(f64, f64)>, String> {
        self.thresholds
            .iter()
            .map(|&threshold| {
                learnable_weight_coordinate_domain(threshold)?.ok_or_else(|| {
                    "smooth-threshold cannot learn a zero threshold multiplicatively".to_string()
                })
            })
            .collect()
    }

    fn value(&self, target: ArrayView1<'_, f64>, rho: ArrayView1<'_, f64>) -> f64 {
        let d = self.latent_dim;
        let n_obs = target.len() / d;
        let mut acc = 0.0;
        for row in 0..n_obs {
            let base = row * d;
            for axis in 0..d {
                let tau = self.threshold(axis, rho);
                let gate = self.sigmoid_gate((target[base + axis] - tau) / self.smoothing_eps);
                acc += self.weight * tau * gate;
            }
        }
        acc
    }

    fn grad_target(&self, target: ArrayView1<'_, f64>, rho: ArrayView1<'_, f64>) -> Array1<f64> {
        let d = self.latent_dim;
        let n_obs = target.len() / d;
        let mut grad = Array1::<f64>::zeros(target.len());
        for row in 0..n_obs {
            let base = row * d;
            for axis in 0..d {
                let tau = self.threshold(axis, rho);
                let gate = self.sigmoid_gate((target[base + axis] - tau) / self.smoothing_eps);
                grad[base + axis] = self.weight * tau * gate * (1.0 - gate) / self.smoothing_eps;
            }
        }
        grad
    }

    fn hessian_diag(
        &self,
        target: ArrayView1<'_, f64>,
        rho: ArrayView1<'_, f64>,
    ) -> Option<Array1<f64>> {
        let d = self.latent_dim;
        let n_obs = target.len() / d;
        let mut diag = Array1::<f64>::zeros(target.len());
        for row in 0..n_obs {
            let base = row * d;
            for axis in 0..d {
                let tau = self.threshold(axis, rho);
                let gate = self.sigmoid_gate((target[base + axis] - tau) / self.smoothing_eps);
                diag[base + axis] = self.true_hessian_diag_entry(tau, gate);
            }
        }
        Some(diag)
    }

    fn hvp(
        &self,
        target: ArrayView1<'_, f64>,
        rho: ArrayView1<'_, f64>,
        v: ArrayView1<'_, f64>,
    ) -> Array1<f64> {
        assert_eq!(target.len(), v.len(), "hvp dimension mismatch");
        let d = self.latent_dim;
        let n_obs = target.len() / d;
        let mut out = Array1::<f64>::zeros(target.len());
        for row in 0..n_obs {
            let base = row * d;
            for axis in 0..d {
                let tau = self.threshold(axis, rho);
                let gate = self.sigmoid_gate((target[base + axis] - tau) / self.smoothing_eps);
                out[base + axis] = self.true_hessian_diag_entry(tau, gate) * v[base + axis];
            }
        }
        out
    }

    fn psd_majorizer_diag(
        &self,
        target: ArrayView1<'_, f64>,
        rho: ArrayView1<'_, f64>,
    ) -> Option<Array1<f64>> {
        // The smooth threshold penalty's exact diagonal Hessian
        //   λτ·g(1−g)(1−2g)/ε²
        // is indefinite (negative once the gate passes the inflection
        // g = ½). The Newton / PIRLS pipeline needs a PSD curvature block, so
        // expose the PSD upper bound implemented by `psd_hessian_diag_entry`:
        // the elementwise max of the re-weighted surrogate and the absolute
        // exact curvature.
        let d = self.latent_dim;
        let n_obs = target.len() / d;
        let mut diag = Array1::<f64>::zeros(target.len());
        for row in 0..n_obs {
            let base = row * d;
            for axis in 0..d {
                let tau = self.threshold(axis, rho);
                let gate = self.sigmoid_gate((target[base + axis] - tau) / self.smoothing_eps);
                diag[base + axis] = self.psd_hessian_diag_entry(tau, gate);
            }
        }
        Some(diag)
    }

    fn grad_rho(&self, target: ArrayView1<'_, f64>, rho: ArrayView1<'_, f64>) -> Array1<f64> {
        let d = self.latent_dim;
        let n_obs = target.len() / d;
        let mut out = Array1::<f64>::zeros(d);
        for axis in 0..d {
            let tau = self.threshold(axis, rho);
            let mut g_tau = 0.0;
            for row in 0..n_obs {
                let x = target[row * d + axis];
                let gate = self.sigmoid_gate((x - tau) / self.smoothing_eps);
                g_tau += gate - tau * gate * (1.0 - gate) / self.smoothing_eps;
            }
            out[axis] = self.weight * tau * g_tau;
        }
        out
    }

    fn rho_count(&self) -> usize {
        self.latent_dim
    }

    fn name(&self) -> &str {
        "smooth_threshold"
    }

    impl_scalar_apply_schedule!(weight);
}

#[cfg(test)]
mod fisher_majorizer_1419_tests {
    use super::*;
    use approx::assert_abs_diff_eq;
    use gam_linalg::faer_ndarray::FaerEigh;
    use ndarray::Array2;

    /// #1419 — the Fisher information metric `G = scale·(diag(a) − a aᵀ)` is PSD
    /// but is NOT a curvature majorizer of the exact softmax-entropy Hessian
    /// `H_entropy`: `G − H_entropy` is indefinite. The genuine Gershgorin
    /// diagonal operator `D_kk = Σ_j|H_kj|` (now `row_psd_majorizer`) IS a
    /// Loewner majorizer: `D − H_entropy ⪰ 0` AND `D ⪰ 0`.
    ///
    /// Oracle: the exact entropy Hessian is built independently from
    /// `row_dense_hessian` (the formula at sparsity.rs:160-193); the smallest
    /// eigenvalue of `M − H` is computed by a direct symmetric eigensolve. The
    /// stated K=2 counterexample (`a=(0.95,0.05)`, `λ=τ=1`) is pinned numerically
    /// against the issue's `H_11 = 0.0783747664` and `G_11 = 0.0475`, and the
    /// contrast (Fisher FAILS, Gershgorin PASSES) is asserted in both the full
    /// K×K block and the single free direction of the reference-logit chart.
    #[test]
    fn gershgorin_majorizes_entropy_where_fisher_does_not_1419() {
        // K=2, λ=τ=1 ⇒ scale = λ/τ² = 1. Logits that realize a = (0.95, 0.05):
        // softmax([z0,z1]) = (0.95,0.05) ⟹ z0 − z1 = ln(0.95/0.05) = ln(19).
        let temperature = 1.0_f64;
        let scale = 1.0_f64; // λ/τ² with λ=1, τ=1.
        let pen = SoftmaxAssignmentSparsityPenalty::new(2, temperature);
        let z1 = 0.0_f64;
        let z0 = z1 + (0.95_f64 / 0.05_f64).ln();
        let row = [z0, z1];

        // Confirm the realized softmax weights.
        let a = pen.softmax_row(&row);
        assert_abs_diff_eq!(a[0], 0.95, epsilon = 1e-12);
        assert_abs_diff_eq!(a[1], 0.05, epsilon = 1e-12);

        // Independent oracles: exact entropy Hessian, Fisher metric, majorizer.
        let h = pen.row_dense_hessian(&row, scale);
        let g = pen.row_fisher_metric(&row, scale);
        let m = pen.row_psd_majorizer(&row, scale);

        // Pin the issue's exact numbers in the sole free direction (index 0):
        //   H_11 = 0.0783747664,  G_11 = a0·a1 = 0.0475.
        assert_abs_diff_eq!(h[[0, 0]], 0.0783747664, epsilon = 1e-9);
        assert_abs_diff_eq!(g[[0, 0]], 0.95 * 0.05, epsilon = 1e-12);

        // The genuine majorizer's diagonal is the abs-row-sum D_kk = Σ_j|H_kj|,
        // raised by the #2339 soft-abs envelope: it must DOMINATE the hard radius
        // (that domination is what carries D̃ ⪰ D ⪰ H) and must sit within the
        // derived relative budget of it.
        for kk in 0..2 {
            let row_sum: f64 = (0..2).map(|jj| h[[kk, jj]].abs()).sum();
            assert!(
                m[[kk, kk]] >= row_sum,
                "smooth Gershgorin radius must DOMINATE the hard radius (#2339); \
                 D̃_{kk} = {} < Σ_j|H_{kk}j| = {row_sum}",
                m[[kk, kk]]
            );
            assert!(
                m[[kk, kk]] - row_sum <= SPECTRAL_DEFLATION_REL_FLOOR * row_sum,
                "smooth Gershgorin radius must stay within the deflation-floor \
                 budget of the hard radius (#2339); D̃_{kk} − D_{kk} = {} > {}",
                m[[kk, kk]] - row_sum,
                SPECTRAL_DEFLATION_REL_FLOOR * row_sum
            );
        }
        // M is a nonnegative diagonal (PSD by inspection) — off-diagonals zero.
        assert_abs_diff_eq!(m[[0, 1]], 0.0, epsilon = 1e-15);
        assert_abs_diff_eq!(m[[1, 0]], 0.0, epsilon = 1e-15);
        assert!(m[[0, 0]] >= 0.0 && m[[1, 1]] >= 0.0);

        // Reference-logit chart: hold z1 fixed, the only free direction is z0, so
        // the reduced 1×1 curvature is the (0,0) entry. Fisher FAILS the Loewner
        // bound there (G_11 − H_11 < 0), the Gershgorin majorizer PASSES it.
        let fisher_free = g[[0, 0]] - h[[0, 0]];
        let major_free = m[[0, 0]] - h[[0, 0]];
        assert!(
            fisher_free < -1e-3,
            "Fisher must FAIL the majorizer bound in the free direction (#1419); \
             G_11 − H_11 = {fisher_free}"
        );
        assert!(
            major_free >= -1e-12,
            "Gershgorin majorizer must SATISFY the bound in the free direction (#1419); \
             D_11 − H_11 = {major_free}"
        );

        // Full K×K Loewner check via a direct symmetric eigensolve oracle.
        // smallest eigenvalue of (M − H) ≥ −tiny ⟹ M ⪰ H; the Fisher case has a
        // strictly negative smallest eigenvalue ⟹ G ⋡ H.
        let mut m_minus_h = Array2::<f64>::zeros((2, 2));
        let mut g_minus_h = Array2::<f64>::zeros((2, 2));
        for i in 0..2 {
            for j in 0..2 {
                m_minus_h[[i, j]] = m[[i, j]] - h[[i, j]];
                g_minus_h[[i, j]] = g[[i, j]] - h[[i, j]];
            }
        }
        let (m_evals, _) = m_minus_h.eigh(faer::Side::Lower).expect("eigh(M−H)");
        let (g_evals, _) = g_minus_h.eigh(faer::Side::Lower).expect("eigh(G−H)");
        let m_min = m_evals.iter().cloned().fold(f64::INFINITY, f64::min);
        let g_min = g_evals.iter().cloned().fold(f64::INFINITY, f64::min);
        assert!(
            m_min >= -1e-12,
            "Gershgorin majorizer must be a Loewner majorizer (M − H ⪰ 0, #1419); \
             smallest eigenvalue of M−H = {m_min}"
        );
        assert!(
            g_min < -1e-9,
            "the OLD Fisher metric must FAIL the Loewner majorizer test (#1419); \
             smallest eigenvalue of G−H = {g_min} (expected strictly negative)"
        );
    }

    /// #1419 — the majorizer's θ-derivative is the exact derivative of the
    /// operator the assembly installs, so value and log-det adjoint differentiate
    /// the SAME `D̃`. Oracle: a central finite difference of `row_psd_majorizer`
    /// itself. FD is permitted ONLY inside this test as an independent check of
    /// the closed-form derivative.
    #[test]
    fn gershgorin_majorizer_logit_derivative_matches_fd_1419() {
        let pen = SoftmaxAssignmentSparsityPenalty::new(4, 0.8);
        let row = [0.3_f64, -0.6, 0.9, 0.2];
        let scale = 1.1_f64 * (1.0 / 0.8_f64) * (1.0 / 0.8_f64);
        let eps = 1e-6;
        for w in 0..4 {
            let dd = pen.row_psd_majorizer_logit_derivative(&row, scale, w);
            let mut rp = row;
            let mut rm = row;
            rp[w] += eps;
            rm[w] -= eps;
            let mp = pen.row_psd_majorizer(&rp, scale);
            let mm = pen.row_psd_majorizer(&rm, scale);
            for k in 0..4 {
                let fd = (mp[[k, k]] - mm[[k, k]]) / (2.0 * eps);
                assert_abs_diff_eq!(dd[[k, k]], fd, epsilon = 1e-6);
            }
            // The derivative is a pure diagonal (D is diagonal).
            for i in 0..4 {
                for j in 0..4 {
                    if i != j {
                        assert_abs_diff_eq!(dd[[i, j]], 0.0, epsilon = 1e-15);
                    }
                }
            }
        }
    }
}

#[cfg(test)]
mod soft_abs_gershgorin_2339_tests {
    //! #2339 (Gershgorin half of #2337 step 1) — the Gershgorin curvature bound
    //! `D_kk = Σ_j|H_kj|` is replaced by the soft-abs envelope
    //! `D̃_kk = Σ_j sqrt(H_kj² + ε₀²‖H_k·‖₂²)`. These gate the four properties the
    //! replacement has to have, each in a form that FAILS if the property is lost:
    //!
    //! 1. MAJORIZATION — `σ_ε ≥ |·|` entrywise and `D̃ ⪰ D ⪰ H`, `D̃ ⪰ 0`. A
    //!    smoothing that dips below `|x|` (the popular `x·tanh(x/ε)` /
    //!    `ε·ln cosh(x/ε)` forms do) breaks the Loewner bound the assembled
    //!    evidence block depends on; the first test pins both directions.
    //! 2. SMOOTHNESS — the θ-adjoint is continuous across a zero crossing of an
    //!    off-diagonal, where the hard `sign(H_kj)` jumps by `2|Ḣ_kj|`. The test
    //!    measures BOTH so the smooth bound cannot pass vacuously.
    //! 3. TIGHTNESS — `0 ≤ D̃_kk − D_kk ≤ SPECTRAL_DEFLATION_REL_FLOOR·D_kk`, the
    //!    derived gap, checked on rows that actually straddle a crossing (where
    //!    the gap is largest and strictly positive).
    //! 4. SCALE DERIVATION — smoothing at the row's OWN `‖H_k·‖₂` keeps `D̃`
    //!    exactly degree-one homogeneous in `scale = λ/τ²`, which is what keeps
    //!    `∂B/∂ρ_sparse` on its existing seam. A fixed absolute `ε` would fail
    //!    this test.
    //!
    //! Finite differences appear ONLY here, as an independent oracle for the
    //! hand-derived closed forms.
    use super::*;
    use approx::assert_abs_diff_eq;
    use gam_linalg::faer_ndarray::FaerEigh;
    use gam_linalg::utils::splitmix64;
    use ndarray::Array2;

    /// Deterministic logit rows spanning the regimes the majorizer sees: a
    /// near-uniform row (where the entropy Hessian is indefinite and the
    /// majorizer earns its keep), a sharply peaked row, and seeded pseudo-random
    /// rows. `splitmix64` keeps this reproducible without a RNG dependency.
    fn seeded_rows(k: usize, seed: u64) -> Vec<Vec<f64>> {
        let mut state = seed;
        let mut rows = vec![vec![0.02_f64; k], {
            let mut peaked = vec![-4.5_f64; k];
            peaked[0] = 3.0;
            peaked[k / 2] = 2.25;
            peaked
        }];
        for _ in 0..6 {
            let row: Vec<f64> = (0..k)
                .map(|_| {
                    let bits = splitmix64(&mut state) >> 11;
                    (bits as f64) / ((1_u64 << 53) as f64) * 8.0 - 4.0
                })
                .collect();
            rows.push(row);
        }
        rows
    }

    /// Hard Gershgorin radius `Σ_j|H_kj|` accumulated in the SAME diagonal-first,
    /// then `j ≠ k` ascending order the smooth radius uses. Same order matters:
    /// `f64` addition is monotone in each addend, so term-wise domination
    /// (`σ_ε ≥ |·|`) implies `D̃_kk ≥ D_kk` EXACTLY only when both sums are
    /// accumulated identically. Comparing against a differently-ordered sum would
    /// weaken a hard guarantee into an approximate one.
    fn hard_radius(h: &Array2<f64>, kk: usize, k: usize) -> f64 {
        let mut acc = h[[kk, kk]].abs();
        for jj in 0..k {
            if jj != kk {
                acc += h[[kk, jj]].abs();
            }
        }
        acc
    }

    /// The OLD, non-smooth θ-derivative `∂D_kk/∂z_w = Σ_j sign(H_kj)·Ḣ_kj` with
    /// `sign(0) = 0`, kept here as the independent reference the smooth adjoint is
    /// contrasted against. It is what `row_psd_majorizer_logit_derivative`
    /// computed before #2339.
    fn hard_radius_logit_derivative(
        pen: &SoftmaxAssignmentSparsityPenalty,
        row: &[f64],
        scale: f64,
        w: usize,
        kk: usize,
    ) -> f64 {
        let k = pen.k_atoms;
        let h = pen.row_dense_hessian(row, scale);
        let dh = pen.row_dense_hessian_logit_derivative(row, scale, w);
        let mut acc = 0.0_f64;
        for jj in 0..k {
            if h[[kk, jj]] != 0.0 {
                acc += h[[kk, jj]].signum() * dh[[kk, jj]];
            }
        }
        acc
    }

    /// Logit `z_0` at which the off-diagonal `H_01` crosses zero — the kink site
    /// of the hard radius — with the remaining logits held at `base`. Located by
    /// a scan for a sign change followed by bisection on the PRODUCTION
    /// `row_dense_hessian`, so the fixture cannot drift from the operator it
    /// probes.
    ///
    /// `K = 2` CANNOT serve as this fixture. The entropy block is gauge-null
    /// (`Σ_j H_kj = 0`, the softmax's shift invariance), so with two atoms row 0
    /// is exactly `(x, −x)`: both entries vanish together and the row DEGENERATES
    /// at the crossing instead of exposing an isolated one. `‖H_0·‖₂` then
    /// collapses to ~1e-16, the smoothing scale `ε₀‖H_0·‖₂` with it, and every
    /// seam probe lands below the representable neighbourhood of `z*`. `K = 3`
    /// keeps `H_00 = −H_02 ≠ 0` where `H_01 = 0`, which is the configuration the
    /// production assembly actually sees.
    fn off_diagonal_zero_crossing(
        pen: &SoftmaxAssignmentSparsityPenalty,
        base: &[f64],
        scale: f64,
    ) -> f64 {
        assert!(
            pen.k_atoms >= 3,
            "the isolated-crossing fixture needs K ≥ 3 (gauge-null rows make K=2 \
             degenerate); got K = {}",
            pen.k_atoms
        );
        let entry = |z0: f64| {
            let mut row = base.to_vec();
            row[0] = z0;
            pen.row_dense_hessian(&row, scale)[[0, 1]]
        };
        let (sweep_lo, sweep_hi, steps) = (-8.0_f64, 8.0_f64, 1600_usize);
        let mut bracket: Option<(f64, f64)> = None;
        let mut prev_z = sweep_lo;
        let mut prev = entry(prev_z);
        for i in 1..=steps {
            let z = sweep_lo + (sweep_hi - sweep_lo) * (i as f64) / (steps as f64);
            let cur = entry(z);
            if prev * cur < 0.0 {
                bracket = Some((prev_z, z));
                break;
            }
            prev_z = z;
            prev = cur;
        }
        let (mut lo, mut hi) =
            bracket.expect("H_01 must change sign over the swept z_0 range (#2339)");
        let lo_is_positive = entry(lo) > 0.0;
        for _ in 0..200 {
            let mid = 0.5 * (lo + hi);
            if mid <= lo || mid >= hi {
                break;
            }
            if (entry(mid) > 0.0) == lo_is_positive {
                lo = mid;
            } else {
                hi = mid;
            }
        }
        0.5 * (lo + hi)
    }

    /// (1) The envelope is an UPPER bound on `|·|` — unconditionally, including
    /// where `f64` rounding of `sqrt(x² + ε²)` would otherwise land below `|x|` —
    /// and exceeds it by at most `ε`. The contrast arm shows the gate is not
    /// vacuous: `x·tanh(x/ε)`, a smooth "soft abs" that is commonly reached for,
    /// sits strictly BELOW `|x|` and would silently invalidate `D ⪰ H`.
    #[test]
    fn soft_abs_envelope_dominates_absolute_value_2339() {
        let mut state = 0x2339_0001_u64;
        let magnitudes = [0.0_f64, 1e-300, 1e-30, 1e-12, 1e-8, 1e-3, 1.0, 7.5, 1e6];
        for &eps in &[0.0_f64, 1e-16, 1e-12, 1e-8, 1e-3, 1.0] {
            let eps_sq = eps * eps;
            for &mag in &magnitudes {
                for sign in [1.0_f64, -1.0] {
                    let x = sign * mag;
                    let env = soft_abs_squared_scale(x, eps_sq);
                    assert!(
                        env >= x.abs(),
                        "soft-abs must MAJORIZE |x| (#2339): σ({x}, ε²={eps_sq}) = {env} \
                         < |x| = {}",
                        x.abs()
                    );
                    assert!(
                        env <= x.abs() + eps + f64::EPSILON * (1.0 + x.abs()),
                        "soft-abs must exceed |x| by at most ε (#2339): \
                         σ({x}, ε²={eps_sq}) − |x| = {} > ε = {eps}",
                        env - x.abs()
                    );
                }
            }
            // At the seam the envelope is EXACTLY ε: strictly above |0| whenever
            // ε > 0, which is precisely the kink fill.
            assert_abs_diff_eq!(
                soft_abs_squared_scale(0.0, eps_sq),
                eps,
                epsilon = 1e-15 * (1.0 + eps)
            );
        }
        // Seeded sweep over arbitrary (x, ε) pairs.
        for _ in 0..4096 {
            let x = (splitmix64(&mut state) >> 11) as f64 / ((1_u64 << 53) as f64) * 20.0 - 10.0;
            let eps = (splitmix64(&mut state) >> 11) as f64 / ((1_u64 << 53) as f64) * 2.0;
            let env = soft_abs_squared_scale(x, eps * eps);
            assert!(
                env >= x.abs() && env <= x.abs() + eps + f64::EPSILON * (1.0 + x.abs()),
                "soft-abs envelope violated at x={x}, ε={eps}: got {env}"
            );
        }
        // Non-vacuity: the smooth alternative that DIPS below |x| is rejected by
        // the same predicate, at the very seam where the difference matters.
        // Sampled inside `tanh`'s transition (it saturates to exactly 1.0 in f64
        // beyond |arg| ≈ 19, where the minorant becomes indistinguishable from
        // |x| and the distinction this arm makes would be invisible).
        let eps = 1e-3_f64;
        for &x in &[1e-4_f64, 1e-3, 5e-3, 1e-2] {
            let dipping = x * (x / eps).tanh();
            assert!(
                dipping < x.abs(),
                "x·tanh(x/ε) is a MINORANT of |x| and must fail the majorization \
                 predicate (#2339): at x={x} it gives {dipping} ≥ |x|"
            );
        }
    }

    /// (1) + (3) On real softmax rows: the smooth radius dominates the hard
    /// radius, the smoothed diagonal is still a Loewner majorizer of the exact
    /// (indefinite) entropy Hessian, and the excess sits inside the derived
    /// deflation-floor budget.
    #[test]
    fn smooth_gershgorin_majorizes_entropy_within_the_derived_budget_2339() {
        let mut checked_rows = 0_usize;
        for (k, temperature, scale) in [
            (2_usize, 1.0_f64, 1.0_f64),
            (3, 0.75, 2.5),
            (5, 1.4, 0.3),
            (8, 0.6, 1.7),
        ] {
            let pen = SoftmaxAssignmentSparsityPenalty::new(k, temperature);
            for row in seeded_rows(k, 0x2339_0000 + k as u64) {
                let h = pen.row_dense_hessian(&row, scale);
                let d_smooth = pen.psd_majorizer_abs_row_sums(&row, scale);
                let mut max_abs_h = 0.0_f64;
                for kk in 0..k {
                    for jj in 0..k {
                        max_abs_h = max_abs_h.max(h[[kk, jj]].abs());
                    }
                }
                for kk in 0..k {
                    let hard = hard_radius(&h, kk, k);
                    assert!(
                        d_smooth[kk] >= hard,
                        "smooth radius must dominate the hard radius EXACTLY \
                         (#2339, k={k}, atom {kk}): {} < {hard}",
                        d_smooth[kk]
                    );
                    assert!(
                        d_smooth[kk] >= 0.0,
                        "the majorizer diagonal must stay nonnegative so D̃ ⪰ 0 \
                         (#2339, k={k}, atom {kk}): {}",
                        d_smooth[kk]
                    );
                    assert!(
                        d_smooth[kk] - hard <= SPECTRAL_DEFLATION_REL_FLOOR * hard,
                        "smoothing gap must stay inside the derived relative budget \
                         SPECTRAL_DEFLATION_REL_FLOOR·D_kk (#2339, k={k}, atom {kk}): \
                         gap {} > budget {}",
                        d_smooth[kk] - hard,
                        SPECTRAL_DEFLATION_REL_FLOOR * hard
                    );
                }
                // Loewner: λ_min(D̃ − H) ≥ 0 by direct symmetric eigensolve.
                let mut gap_matrix = Array2::<f64>::zeros((k, k));
                for i in 0..k {
                    for j in 0..k {
                        gap_matrix[[i, j]] = -h[[i, j]];
                    }
                    gap_matrix[[i, i]] += d_smooth[i];
                }
                let (evals, _) = gap_matrix.eigh(faer::Side::Lower).expect("eigh(D̃−H)");
                let min_eig = evals.iter().cloned().fold(f64::INFINITY, f64::min);
                assert!(
                    min_eig >= -1e-12 * (1.0 + max_abs_h),
                    "the SMOOTHED Gershgorin diagonal must remain a Loewner \
                     majorizer of the exact entropy Hessian (#1419/#2339, k={k}): \
                     λ_min(D̃−H) = {min_eig}"
                );
                checked_rows += 1;
            }
        }
        assert!(
            checked_rows == 32,
            "expected 32 gated rows, ran {checked_rows}"
        );
    }

    /// The seam fixture: a `K = 3` logit row sitting exactly on an `H_01` zero
    /// crossing, plus the width of the smoothing band around it — the envelope's
    /// curvature-space scale `ε_0 = ε₀·‖H_0·‖₂` divided by the crossing entry's
    /// slope `|∂H_01/∂z_0|`. Every seam probe below is expressed as a fraction of
    /// THIS band, never as a hard-coded offset: the band is set by a derived
    /// temperature and the fixture's own curvature, so a literal step would be a
    /// magic constant that silently stops probing the seam if either moves.
    /// Returns `(row, band, |Ḣ_01|)`.
    fn zero_crossing_fixture(
        pen: &SoftmaxAssignmentSparsityPenalty,
        base: &[f64],
        scale: f64,
    ) -> (Vec<f64>, f64, f64) {
        let z_star = off_diagonal_zero_crossing(pen, base, scale);
        let mut row = base.to_vec();
        row[0] = z_star;
        let h = pen.row_dense_hessian(&row, scale);
        let norm = (0..pen.k_atoms)
            .map(|jj| h[[0, jj]] * h[[0, jj]])
            .sum::<f64>()
            .sqrt();
        let eps_0 = SoftmaxAssignmentSparsityPenalty::soft_abs_temperature(pen.k_atoms) * norm;
        let slope = pen.row_dense_hessian_logit_derivative(&row, scale, 0)[[0, 1]].abs();
        (row, eps_0 / slope, slope)
    }

    /// (2) The θ-adjoint is CONTINUOUS across an off-diagonal zero crossing,
    /// where the hard `sign(H_kj)` jumps by `2|Ḣ_kj|`. Both sides are measured:
    /// the smooth jump must shrink with the probe offset (C⁰ derivative), while
    /// the hard jump must stay bounded away from zero — otherwise this test would
    /// pass on an operator with no kink to remove.
    #[test]
    fn smooth_gershgorin_adjoint_is_continuous_across_a_zero_crossing_2339() {
        let scale = 1.3_f64;
        let pen = SoftmaxAssignmentSparsityPenalty::new(3, 1.0);
        let base = [0.0_f64, 0.0, -0.7];
        let (row, band, slope) = zero_crossing_fixture(&pen, &base, scale);
        assert!(
            slope > 1e-4 && band > 1e-13 && band.is_finite(),
            "the crossing entry must move under z_0, and the row must not \
             degenerate, for the kink to be real (#2339): |∂H_01/∂z_0| = {slope}, \
             band = {band}"
        );

        let mut measured: Vec<(f64, f64, f64)> = Vec::new();
        for divisor in [100.0_f64, 1000.0] {
            let delta = band / divisor;
            let mut plus = row.clone();
            let mut minus = row.clone();
            plus[0] += delta;
            minus[0] -= delta;
            // The realized offset after rounding into `z*`'s exponent.
            let realized = 0.5 * (plus[0] - minus[0]);
            assert!(
                realized > 0.0,
                "the seam probe must be representable next to z* = {} (#2339): \
                 requested δ = {delta}",
                row[0]
            );
            let smooth = (pen.row_psd_majorizer_logit_derivative(&plus, scale, 0)[[0, 0]]
                - pen.row_psd_majorizer_logit_derivative(&minus, scale, 0)[[0, 0]])
            .abs();
            let hard = (hard_radius_logit_derivative(&pen, &plus, scale, 0, 0)
                - hard_radius_logit_derivative(&pen, &minus, scale, 0, 0))
            .abs();
            measured.push((realized, smooth, hard));
        }

        for &(delta, smooth, hard) in &measured {
            assert!(
                hard >= slope,
                "the HARD radius must genuinely jump by ≈2|Ḣ_01| across the \
                 crossing — otherwise there is no kink for #2339 to remove \
                 (δ={delta}): hard jump {hard} < {slope}"
            );
            assert!(
                smooth <= 0.05 * hard,
                "the SMOOTH adjoint must not jump across the crossing (#2339, \
                 δ={delta}): smooth jump {smooth} vs hard jump {hard}"
            );
        }
        // C⁰: the one-sided values converge as the probe closes in. A ten-fold
        // smaller offset must cut the residual jump by at least four.
        let (delta_coarse, coarse, _) = measured[0];
        let (delta_fine, fine, _) = measured[1];
        assert!(
            fine <= 0.25 * coarse,
            "the smooth adjoint's residual jump must vanish with the probe offset \
             (#2339): {coarse} at δ={delta_coarse} → {fine} at δ={delta_fine}"
        );
    }

    /// (2) The smooth radius is differentiable AT the seam in the strong sense:
    /// a central finite difference taken INSIDE the smoothing band recovers the
    /// hand-derived closed form. With the hard `|·|` this is the #2253 failure
    /// mode — the stencil straddles the kink and the FD reference is meaningless,
    /// off by `O(|Ḣ_01|)`. The tolerance sits ~200× below that failure signal and
    /// ~25× above the FD's own truncation/cancellation floor at a band/100 step
    /// (truncation grows like `h²·|Ḣ_01|³/ε_k²`, cancellation like `ulp(D̃)/h`;
    /// band/100 is near their crossover).
    #[test]
    fn smooth_gershgorin_adjoint_matches_fd_inside_the_smoothing_band_2339() {
        let scale = 1.3_f64;
        let pen = SoftmaxAssignmentSparsityPenalty::new(3, 1.0);
        let base = [0.0_f64, 0.0, -0.7];
        let (row, band, slope) = zero_crossing_fixture(&pen, &base, scale);
        let step = band / 100.0;
        assert!(
            step > 1e-15 && step.is_finite() && slope > 1e-4,
            "the ε-scaled probe step must be usable (#2339): step = {step}, \
             slope = {slope}"
        );
        for w in 0..3 {
            let analytic = pen.row_psd_majorizer_logit_derivative(&row, scale, w);
            let mut plus = row.clone();
            let mut minus = row.clone();
            plus[w] += step;
            minus[w] -= step;
            let realized = 0.5 * (plus[w] - minus[w]);
            let mp = pen.row_psd_majorizer(&plus, scale);
            let mm = pen.row_psd_majorizer(&minus, scale);
            for kk in 0..3 {
                let fd = (mp[[kk, kk]] - mm[[kk, kk]]) / (2.0 * realized);
                assert_abs_diff_eq!(analytic[[kk, kk]], fd, epsilon = 1e-3);
            }
        }
    }

    /// (4) `D̃` is EXACTLY degree-one homogeneous in `scale = λ/τ²` — bit-for-bit
    /// under a power-of-two rescale, and to rounding under an arbitrary one. This
    /// is the property that keeps the `∂B/∂ρ_sparse` channel on its existing seam
    /// without a code change; smoothing at a fixed absolute `ε` instead of the
    /// row's own `‖H_k·‖₂` would break it.
    #[test]
    fn smooth_gershgorin_is_degree_one_homogeneous_in_scale_2339() {
        let k = 5_usize;
        let pen = SoftmaxAssignmentSparsityPenalty::new(k, 0.9);
        let base = 0.625_f64;
        for row in seeded_rows(k, 0x2339_0100) {
            let d_base = pen.psd_majorizer_abs_row_sums(&row, base);
            let d_doubled = pen.psd_majorizer_abs_row_sums(&row, 2.0 * base);
            let d_tilted = pen.psd_majorizer_abs_row_sums(&row, 3.5 * base);
            for kk in 0..k {
                assert_eq!(
                    d_doubled[kk],
                    2.0 * d_base[kk],
                    "D̃ must be degree-1 homogeneous in scale BIT-FOR-BIT under a \
                     power-of-two rescale (#2339, atom {kk})"
                );
                assert_abs_diff_eq!(
                    d_tilted[kk],
                    3.5 * d_base[kk],
                    epsilon = 1e-14 * (1.0 + d_base[kk])
                );
            }
            // The adjoint inherits the same homogeneity (it is ∂ of a degree-1
            // homogeneous operator at fixed logits).
            let dd_base = pen.row_psd_majorizer_logit_derivative(&row, base, 0);
            let dd_doubled = pen.row_psd_majorizer_logit_derivative(&row, 2.0 * base, 0);
            for kk in 0..k {
                assert_eq!(
                    dd_doubled[[kk, kk]],
                    2.0 * dd_base[[kk, kk]],
                    "∂D̃/∂z must be degree-1 homogeneous in scale BIT-FOR-BIT \
                     (#2339, atom {kk})"
                );
            }
        }
    }

    /// (4), the consequence production actually depends on: the majorizer is
    /// **scale-equivariant**, so the two routes by which a #991 per-row design
    /// weight `w_i` reaches it agree.
    ///
    /// The envelope is `σ_ε(x) = ε·h(x/ε)` with `h(u) = sqrt(u²+1)` a fixed smooth
    /// majorizer of `|·|` (worst-case gap `ε·h(0) = ε`, linear in ε), and `ε` is
    /// taken from the operator's OWN magnitude `ε₀‖H_k·‖₂`. Both factors are
    /// degree-1 in `H_k·`, so `D̃(c·H) = c·D̃(H)` for every `c > 0`: a Gershgorin
    /// bound that moved under a rescaling of the model would be a bug in the bound
    /// itself.
    ///
    /// That is not an abstract property here — two production call sites weight
    /// the SAME block differently and rely on the two being equal:
    ///
    /// * `psd_majorizer_diag` (the trait channel) computes `D̃` at the unweighted
    ///   `scale` and post-multiplies: `w_i·D̃(scale)`.
    /// * the SAE arrow-Schur assembly folds the weight into the strength instead:
    ///   `row_psd_majorizer(logits, scale·w_i)`, i.e. `D̃(w_i·scale)`.
    ///
    /// With the hard `|·|` this was trivially true term by term. With a smoothed
    /// radius it is true ONLY because `ε` is derived from the row rather than
    /// fixed: a constant `ε` makes `D̃(w·scale) ≠ w·D̃(scale)`, silently splitting
    /// the assembled `B` from the diagonal the trait reports, and no existing test
    /// compares the two routes (`every_channel_scales_by_w_row_identically` checks
    /// the trait channel against ITSELF reweighted, never against the assembly's
    /// folded form). Exact for a power-of-two weight, tight-relative otherwise.
    #[test]
    fn smooth_gershgorin_weighting_routes_agree_by_scale_equivariance_2339() {
        let k = 6_usize;
        let pen = SoftmaxAssignmentSparsityPenalty::new(k, 1.1);
        let scale = 0.75_f64;
        let mut compared = 0_usize;
        for row in seeded_rows(k, 0x2339_0200) {
            // Power-of-two weights: equivariance must hold BIT-FOR-BIT, because
            // scaling by a power of two is exact in every intermediate product.
            for &w in &[0.5_f64, 2.0, 4.0] {
                let post_multiplied = pen.psd_majorizer_abs_row_sums(&row, scale);
                let folded = pen.psd_majorizer_abs_row_sums(&row, scale * w);
                for kk in 0..k {
                    assert_eq!(
                        folded[kk],
                        w * post_multiplied[kk],
                        "the assembly's folded weight D̃(w·scale) and the trait \
                         channel's w·D̃(scale) must be the SAME operator (#991/#2339, \
                         w={w}, atom {kk}); a fixed absolute ε would split them"
                    );
                    compared += 1;
                }
            }
            // Arbitrary weights: equal to rounding. A constant-ε envelope would
            // miss by ≈(w−1)·K·ε, i.e. ~1e-8 relative — far outside this bound.
            for &w in &[0.37_f64, 1.9, 6.25] {
                let post_multiplied = pen.psd_majorizer_abs_row_sums(&row, scale);
                let folded = pen.psd_majorizer_abs_row_sums(&row, scale * w);
                for kk in 0..k {
                    assert_abs_diff_eq!(
                        folded[kk],
                        w * post_multiplied[kk],
                        epsilon = 1e-13 * (1.0 + w * post_multiplied[kk])
                    );
                }
            }
            // The adjoint the assembly pairs with that block carries the weight
            // identically, so value and adjoint cannot drift apart under weighting.
            for &w in &[0.5_f64, 2.0] {
                for probe in 0..k {
                    let post_multiplied =
                        pen.row_psd_majorizer_logit_derivative(&row, scale, probe);
                    let folded = pen.row_psd_majorizer_logit_derivative(&row, scale * w, probe);
                    for kk in 0..k {
                        assert_eq!(
                            folded[[kk, kk]],
                            w * post_multiplied[[kk, kk]],
                            "∂D̃/∂z_w must carry the #991 row weight identically to \
                             the value (#991/#2339, w={w}, atom {kk}, probe {probe})"
                        );
                    }
                }
            }
        }
        assert!(
            compared == 8 * 3 * k,
            "expected {} weighted comparisons, made {compared}",
            8 * 3 * k
        );
    }

    /// (2), degenerate corner: an atom whose softmax mass underflows to exactly
    /// zero has an identically zero Hessian row, so the envelope's smoothing
    /// scale is exactly zero too. Value and adjoint must both be exactly `0.0` —
    /// the same exact-zero continuation `entropy_log_plus_one` uses — rather than
    /// a `0/0` NaN from the soft-sign division.
    #[test]
    fn smooth_gershgorin_is_exactly_zero_on_an_underflowed_atom_2339() {
        let pen = SoftmaxAssignmentSparsityPenalty::new(2, 1.0);
        let scale = 2.0_f64;
        // exp(-800) underflows to exactly 0, so a_1 == 0.0 and H_1· ≡ 0.
        let row = [0.0_f64, -800.0];
        let d = pen.psd_majorizer_abs_row_sums(&row, scale);
        assert_eq!(
            d[1], 0.0,
            "an underflowed atom's majorizer diagonal must be exactly zero (#2339)"
        );
        assert!(
            d[0].is_finite() && d[0] >= 0.0,
            "the surviving atom must stay finite and nonnegative (#2339): {}",
            d[0]
        );
        for w in 0..2 {
            let dd = pen.row_psd_majorizer_logit_derivative(&row, scale, w);
            assert_eq!(
                dd[[1, 1]],
                0.0,
                "an underflowed atom's majorizer adjoint must be exactly zero, \
                 not NaN (#2339, w={w})"
            );
            assert!(
                dd[[0, 0]].is_finite(),
                "the surviving atom's adjoint must stay finite (#2339, w={w}): {}",
                dd[[0, 0]]
            );
        }
    }
}

#[cfg(test)]
mod row_weighted_prior_991_tests {
    //! #991 design-honesty per-row weights: row `i`'s softmax-entropy prior must
    //! be scaled by `w_i` IDENTICALLY in every channel. Because value, gradient,
    //! Hessian diagonal, HVP, and the PSD majorizer are all linear in the per-row
    //! penalty strength, scaling the strength by `w_i` scales all of them by the
    //! same `w_i` and cannot desync them. These are the CI gate for that
    //! invariant (the fit that consumes it cannot be run here).
    use super::AnalyticPenalty;
    use super::*;
    use approx::assert_abs_diff_eq;
    use ndarray::{Array1, s};

    fn logits(n: usize, k: usize) -> Array1<f64> {
        // Deterministic non-uniform logits so every row has genuine entropy
        // gradient/curvature (no trivially-degenerate softmax rows).
        let mut v = Array1::<f64>::zeros(n * k);
        for r in 0..n {
            for a in 0..k {
                v[r * k + a] =
                    0.35 * (r as f64) - 0.6 * (a as f64) + 0.11 * ((r * k + a) as f64).sin();
            }
        }
        v
    }

    /// The weighted value equals the unweighted per-row entropies recombined with
    /// `w_i`, and the mean-1 weighting leaves the total exactly invariant when the
    /// weights average to one — the design-honesty contract.
    #[test]
    fn weighted_value_is_per_row_reweight_of_unweighted() {
        let (n, k) = (5usize, 3usize);
        let temperature = 0.7_f64;
        let rho = Array1::from_vec(vec![0.2_f64]);
        let target = logits(n, k);
        let base = SoftmaxAssignmentSparsityPenalty::new(k, temperature);
        // Per-row entropies via single-row penalties (each a 1-row problem).
        let mut per_row = vec![0.0_f64; n];
        for r in 0..n {
            let row = target.slice(s![r * k..r * k + k]).to_owned();
            per_row[r] = base.value(row.view(), rho.view());
        }
        let unweighted: f64 = per_row.iter().sum();
        assert_abs_diff_eq!(
            base.value(target.view(), rho.view()),
            unweighted,
            epsilon = 1e-12
        );

        let w = vec![1.7_f64, 0.3, 1.1, 0.5, 1.4]; // mean = 1.0 exactly.
        let weighted = base.clone().with_row_weights(Some(&w));
        let expect: f64 = (0..n).map(|r| w[r] * per_row[r]).sum();
        assert_abs_diff_eq!(
            weighted.value(target.view(), rho.view()),
            expect,
            epsilon = 1e-12
        );
        // Mean-1 weights preserve the total (Σ w_i H_i vs Σ H_i differ only by the
        // per-row redistribution, but here we assert the exact reweighted target).
        assert_abs_diff_eq!(
            weighted.value(target.view(), rho.view()),
            (0..n).map(|r| w[r] * per_row[r]).sum::<f64>(),
            epsilon = 1e-12
        );
    }

    /// FD ORACLE: `d(value)/d(z_{r,a}) == grad_target[r*K+a]` under NONTRIVIAL
    /// per-row weights. This is the value/gradient desync gate — if any channel
    /// carried a different weighting than the value, this central difference would
    /// diverge from the analytic gradient.
    #[test]
    fn weighted_value_grad_are_fd_consistent() {
        let (n, k) = (4usize, 3usize);
        let temperature = 0.9_f64;
        let rho = Array1::from_vec(vec![-0.1_f64]);
        let target = logits(n, k);
        let w = vec![1.9_f64, 0.4, 0.8, 0.9];
        let pen = SoftmaxAssignmentSparsityPenalty::new(k, temperature).with_row_weights(Some(&w));
        let grad = pen.grad_target(target.view(), rho.view());
        let eps = 1e-6;
        for idx in 0..n * k {
            let mut plus = target.clone();
            let mut minus = target.clone();
            plus[idx] += eps;
            minus[idx] -= eps;
            let fd = (pen.value(plus.view(), rho.view()) - pen.value(minus.view(), rho.view()))
                / (2.0 * eps);
            assert_abs_diff_eq!(grad[idx], fd, epsilon = 1e-7);
        }
    }

    /// Every channel scales by exactly `w_i` on row `i` relative to the unweighted
    /// penalty — grad_target, hessian_diag, psd_majorizer_diag, and hvp. Confirms
    /// the single strength multiplier reaches all of them identically.
    #[test]
    fn every_channel_scales_by_w_row_identically() {
        let (n, k) = (4usize, 3usize);
        let temperature = 0.8_f64;
        let rho = Array1::from_vec(vec![0.15_f64]);
        let target = logits(n, k);
        let v = logits(n, k); // arbitrary HVP direction.
        let w = vec![1.6_f64, 0.25, 1.05, 1.1];
        let base = SoftmaxAssignmentSparsityPenalty::new(k, temperature);
        let wtd = base.clone().with_row_weights(Some(&w));

        let g0 = base.grad_target(target.view(), rho.view());
        let g1 = wtd.grad_target(target.view(), rho.view());
        let d0 = base.hessian_diag(target.view(), rho.view()).unwrap();
        let d1 = wtd.hessian_diag(target.view(), rho.view()).unwrap();
        let m0 = base.psd_majorizer_diag(target.view(), rho.view()).unwrap();
        let m1 = wtd.psd_majorizer_diag(target.view(), rho.view()).unwrap();
        let h0 = base.hvp(target.view(), rho.view(), v.view());
        let h1 = wtd.hvp(target.view(), rho.view(), v.view());
        for r in 0..n {
            for a in 0..k {
                let i = r * k + a;
                assert_abs_diff_eq!(g1[i], w[r] * g0[i], epsilon = 1e-12);
                assert_abs_diff_eq!(d1[i], w[r] * d0[i], epsilon = 1e-12);
                assert_abs_diff_eq!(m1[i], w[r] * m0[i], epsilon = 1e-12);
                assert_abs_diff_eq!(h1[i], w[r] * h0[i], epsilon = 1e-12);
            }
        }
        // grad_rho (softmax) is the value itself, so it too carries the weighting.
        let r0 = base.grad_rho(target.view(), rho.view())[0];
        let r1 = wtd.grad_rho(target.view(), rho.view())[0];
        let expect: f64 = (0..n)
            .map(|r| {
                let row = target.slice(s![r * k..r * k + k]).to_owned();
                w[r] * base.value(row.view(), rho.view())
            })
            .sum();
        assert_abs_diff_eq!(r1, expect, epsilon = 1e-12);
        assert!(r0.is_finite());
    }

    /// `None` weights are byte-for-byte the unweighted path (no silent ×1.0 drift).
    #[test]
    fn none_weights_are_bit_for_bit_unweighted() {
        let (n, k) = (3usize, 4usize);
        let rho = Array1::from_vec(vec![0.0_f64]);
        let target = logits(n, k);
        let base = SoftmaxAssignmentSparsityPenalty::new(k, 1.0);
        let none = base.clone().with_row_weights(None);
        assert_eq!(
            base.value(target.view(), rho.view()).to_bits(),
            none.value(target.view(), rho.view()).to_bits()
        );
        let g0 = base.grad_target(target.view(), rho.view());
        let g1 = none.grad_target(target.view(), rho.view());
        for i in 0..n * k {
            assert_eq!(g0[i].to_bits(), g1[i].to_bits());
        }
    }
}