Struct SaeAssignment 

pub struct SaeAssignment {
    pub logits: Array2<f64>,
    pub coords: Vec<LatentCoordValues>,
    pub mode: AssignmentMode,
    pub ungated: Vec<bool>,
    pub frozen_logits: Option<Array2<f64>>,
}

Per-row latent assignment state.

The stored assignment parameter is logits; non-negative assignments are derived by row-wise softmax, independent IBP-MAP sigmoid active indicators, or JumpReLU gates. Softmax logits are canonicalized to the reference chart logits[K - 1] = 0, so the row-local Newton coordinates contain only the first K - 1 logits (0 coordinates for K = 1). Gate-style modes keep all K logits as identifiable scalar parameters. coords[k] holds t_{.,k} for atom k.

Fields

logits: Array2<f64>

coords: Vec<LatentCoordValues>

mode: AssignmentMode

ungated: Vec<bool>

#1026 — per-atom UNGATED flag (length K, default all-false). An ungated atom is the dense linear/background tier: its per-row gate is fixed at a_k ≡ 1 (it contributes γ_k(t_k) to EVERY row, unweighted), it is excluded from the other atoms’ gate (for the column-separable IBP / JumpReLU modes the remaining atoms are computed independently, so they are unaffected), and its logit is NOT a free parameter — its logit-JVP, sparsity-prior gradient/curvature, and softmax majorizer contributions are all zero, leaving its logit slot an inert (ridge-regularized) null direction in the per-row Newton block. This lets the linear tier carry FULL-RANK reconstructible variance (fitted = γ_ungated(x) + Σ_{gated} a_k·γ_k(x)) so a linear SAE can reach the rank-(K·d) PCA ceiling, while the gated curved atoms still add sparse structure on the residual (#1026 routing-bound finding).

frozen_logits: Option<Array2<f64>>

#1033 — AMORTIZED / FROZEN routing. When Some, this (n, K) matrix is a ρ-INVARIANT predicted routing (the amortized x → logits map distilled from the frozen dictionary): the gates are computed from THESE logits instead of the free self.logits, and the logits are NOT optimized by the inner Newton (their gradient/curvature/prior contributions are zeroed, exactly as for Self::ungated). This is the generalization of an ungated atom from “pin the gate at 1” to “pin the gate at the predicted value”: it makes the per-row routing a fixed function of x + the frozen dictionary, so the outer ρ-search reuses ONE routing instead of re-solving per-row gates every outer eval — the n-independent-outer-loop lever (#1033). None is the historical free-logit path (bit-identical).

Implementations

impl SaeAssignment

pub fn new( logits: Array2<f64>, coords: Vec<LatentCoordValues>, temperature: f64, ) -> Result<Self, String>

pub fn with_mode( logits: Array2<f64>, coords: Vec<LatentCoordValues>, mode: AssignmentMode, ) -> Result<Self, String>

pub fn with_frozen_routing( self, predicted: Option<Array2<f64>>, ) -> Result<Self, String>

#1033 — install a ρ-INVARIANT FROZEN routing (the amortized predicted logits; see SaeAssignment::frozen_logits). predicted must be (n, K). With routing frozen, the gates are computed from predicted and the logits are excluded from the inner Newton (their gradient/curvature are inert, like an ungated atom’s). Passing None restores the free-logit path.

pub fn routing_is_frozen(&self) -> bool

Whether the per-row routing is FROZEN (amortized) rather than free-logit.

pub fn freeze_routing_from_current_logits(self) -> Result<Self, String>

#1033 — install the simplest faithful AMORTIZED routing predictor: a fixed-form DISTILL of the current dictionary’s routing, namely the current (converged) logits SNAPSHOTTED as the ρ-invariant frozen routing. This is the x → logits map “evaluated once at the frozen dictionary” — the routing the dictionary already expresses — held fixed so the outer ρ-search reuses it instead of re-optimizing the gates at every ρ. (A richer predictor that recomputes logits from x via the encode-atlas chart geometry is a later refinement; snapshotting the converged routing is the exact fixed-point it would target at the frozen dictionary.) Rejected for Softmax for the same simplex-coupling reason as Self::with_frozen_routing.

pub fn freeze_routing_in_place(&mut self) -> Result<(), String>

#1033 — in-place variant of Self::freeze_routing_from_current_logits for callers holding &mut SaeAssignment (e.g. inside a SaeManifoldTerm), where moving the assignment out is awkward. Same contract: snapshot the current logits as the ρ-invariant frozen routing; reject Softmax.

pub fn set_frozen_routing_in_place( &mut self, predicted: Array2<f64>, ) -> Result<(), String>

#1033 — install an explicit predicted routing in place (the RoutingPredictor::ChartGeometry output), &mut self variant of Self::with_frozen_routing. predicted must be (n, K); rejects Softmax (separable-mode contract) and non-finite predictions.

pub fn thaw_routing(&mut self)

#1033 — lift the frozen routing, restoring the free-logit search path.

pub fn with_ungated(self, flags: Vec<bool>) -> Result<Self, String>

#1026 — designate which atoms are UNGATED (the dense linear/background tier; see SaeAssignment::ungated). flags must have length K.

Ungating is defined for the COLUMN-SEPARABLE gate modes (IBP-MAP and JumpReLU): each atom’s gate is an independent per-atom function of its own logit, so pinning one atom to a_k ≡ 1 leaves every other atom’s gate exactly as computed. Softmax is a coupled simplex (Σ_k a_k = 1 over all K), so a unit gate for one atom is only well defined relative to a gated-subset renormalization that must also be reflected in the logit-JVP and the entropy majorizer; this constructor’s contract is restricted to the separable modes, and an ungated atom under Softmax is REJECTED here so the inner solve never runs on a value/gradient-mismatched gate. Callers wanting a dense background tier under Softmax route it as an IBP-MAP or JumpReLU atom.

pub fn has_ungated(&self) -> bool

Whether any atom is ungated (the #1026 background tier is engaged).

pub fn n_obs(&self) -> usize

pub fn k_atoms(&self) -> usize

pub fn total_coord_dim(&self) -> usize

pub fn assignment_coord_dim(&self) -> usize

pub fn row_block_dim(&self) -> usize

pub fn coord_offsets(&self) -> Vec<usize>

pub fn assignments(&self) -> Array2<f64>

pub fn assignments_row(&self, row: usize) -> Array1<f64>

pub fn try_assignments_row(&self, row: usize) -> Result<Array1<f64>, String>

pub fn flatten_ext_coords(&self) -> Array1<f64>

Flatten extension coordinates in row-major SAE layout: (assignment chart_i, t_i0[0..d_0], ..., t_iK[0..d_K]) for every row. Softmax contributes the first K - 1 reference logits and omits the fixed reference logit; gate-style assignment modes contribute all K logits.

pub fn from_blocks_with_mode( logits: Array2<f64>, coord_blocks: Vec<Array2<f64>>, mode: AssignmentMode, ) -> Result<Self, String>

pub fn from_blocks_with_mode_and_manifolds( logits: Array2<f64>, coord_blocks: Vec<Array2<f64>>, manifolds: Vec<LatentManifold>, mode: AssignmentMode, ) -> Result<Self, String>

Trait Implementations

impl Clone for SaeAssignment

fn clone(&self) -> SaeAssignment

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for SaeAssignment

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.