gam-terms 0.3.157

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
    }

    /// 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
    }

}

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 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::utils::splitmix64;

    /// (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|"
            );
        }
    }

}

#[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());
        }
    }
}