pub struct SaeRowLayout {
pub active_atoms: Vec<Vec<usize>>,
pub coord_starts: Vec<Vec<usize>>,
pub coord_offsets_full: Vec<usize>,
pub coord_dims: Vec<usize>,
}Expand description
Per-row active-set layout for sparse SAE assignment (any mode).
When the assignment is sparse — structurally (JumpReLU gate) or
effectively (softmax / IBP-MAP at large K, where the assignment mass
concentrates on a small support) — only a subset of K atoms are active
per observation. The Arrow-Schur row block for observation i has dim
q_active_i = |active_atoms_i| + Σ_{k ∈ active_i} d_k rather than
q = assignment_dim + Σ_k d_k. This struct records which atoms are active per row
and maps compressed block positions back to full-q positions so that
apply_newton_step can unpack the compact delta_t from the solve.
For JumpReLU the active set is exactly the gated support
(a_{n,k} ≠ 0), so the compact solve is identity to the dense solve.
For IBP-MAP the active set is the union of a top-k_active_cap
truncation and a magnitude cutoff on a_{n,k}; this is only enabled when
K is large enough that the dense (m_total · p)² data Gram would not
fit the host / device working-set budget, and the dropped atoms carry
O(a_{n,k}²) curvature that is negligible by construction of the cutoff.
#1408: SOFTMAX engages this compact layout when an explicit top_k
(softmax_active_cap) and/or the in-core memory budget bounds the active
set — the AssignmentMode::Softmax arm of assemble_arrow_schur consults
crate::manifold::SaeManifoldTerm::softmax_active_plan and,
on Some((cap, cutoff)), builds the active set via
Self::from_dense_weights. The full-K dense softmax layout is retained
only when neither lever engages (no top_k, in-budget K). Folding softmax
top_k into the compact solve required writing the active×active Gershgorin
Loewner majorizer sub-block (#1419; the softmax entropy curvature is
indefinite, so its raw diagonal cannot be used) AND contracting that SAME
majorizer over the compact logit slots in the logdet ρ-trace
(assignment_log_strength_hessian_trace) and the θ-adjoint, so value,
log|H|, and Γ differentiate one operator on the compact support. That
coordinated change is landed and FD-certified; the FFI’s after-the-fit
top-k projection is then a no-op at the optimum.
Fields§
§active_atoms: Vec<Vec<usize>>active_atoms[row] — sorted indices of active atoms for that row.
coord_starts: Vec<Vec<usize>>For row i, active atom active_atoms[i][j] has its coord block
starting at compressed position coord_starts[i][j].
coord_offsets_full: Vec<usize>Full-q coordinate offset for atom k (length k_atoms).
coord_dims: Vec<usize>Per-atom coordinate dimensions, indexed by atom index.
Implementations§
Source§impl SaeRowLayout
impl SaeRowLayout
Sourcepub fn row_q_active(&self, row: usize) -> usize
pub fn row_q_active(&self, row: usize) -> usize
Per-row compressed dim.
Sourcepub fn expand_row(&self, row: usize, delta_t_row: &[f64], out: &mut [f64])
pub fn expand_row(&self, row: usize, delta_t_row: &[f64], out: &mut [f64])
Expand a compact delta_t row slice back into full-q, zeros for inactive.
Trait Implementations§
Source§impl Clone for SaeRowLayout
impl Clone for SaeRowLayout
Source§fn clone(&self) -> SaeRowLayout
fn clone(&self) -> SaeRowLayout
1.0.0 (const: unstable) · Source§fn clone_from(&mut self, source: &Self)
fn clone_from(&mut self, source: &Self)
source. Read moreAuto Trait Implementations§
impl Freeze for SaeRowLayout
impl RefUnwindSafe for SaeRowLayout
impl Send for SaeRowLayout
impl Sync for SaeRowLayout
impl Unpin for SaeRowLayout
impl UnsafeUnpin for SaeRowLayout
impl UnwindSafe for SaeRowLayout
Blanket Implementations§
impl<T> Allocation for T
Source§impl<T> BorrowMut<T> for Twhere
T: ?Sized,
impl<T> BorrowMut<T> for Twhere
T: ?Sized,
Source§fn borrow_mut(&mut self) -> &mut T
fn borrow_mut(&mut self) -> &mut T
impl<ST, DT> CastableFrom<ST, Initialized, Initialized> for DT
impl<ST, DT> CastableFrom<ST, Uninit, Uninit> for DT
Source§impl<T> CloneToUninit for Twhere
T: Clone,
impl<T> CloneToUninit for Twhere
T: Clone,
Source§impl<T> DistributionExt for Twhere
T: ?Sized,
impl<T> DistributionExt for Twhere
T: ?Sized,
impl<T, U> Imply<T> for U
Source§impl<T> IntoEither for T
impl<T> IntoEither for T
Source§fn into_either(self, into_left: bool) -> Either<Self, Self>
fn into_either(self, into_left: bool) -> Either<Self, Self>
self into a Left variant of Either<Self, Self>
if into_left is true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read moreSource§fn into_either_with<F>(self, into_left: F) -> Either<Self, Self>
fn into_either_with<F>(self, into_left: F) -> Either<Self, Self>
self into a Left variant of Either<Self, Self>
if into_left(&self) returns true.
Converts self into a Right variant of Either<Self, Self>
otherwise. Read moreSource§impl<T> Pointable for T
impl<T> Pointable for T
impl<T> Read<Exclusive, BecauseExclusive> for Twhere
T: ?Sized,
Source§impl<SS, SP> SupersetOf<SS> for SPwhere
SS: SubsetOf<SP>,
impl<SS, SP> SupersetOf<SS> for SPwhere
SS: SubsetOf<SP>,
Source§fn to_subset(&self) -> Option<SS>
fn to_subset(&self) -> Option<SS>
self from the equivalent element of its
superset. Read moreSource§fn is_in_subset(&self) -> bool
fn is_in_subset(&self) -> bool
self is actually part of its subset T (and can be converted to it).Source§fn to_subset_unchecked(&self) -> SS
fn to_subset_unchecked(&self) -> SS
self.to_subset but without any property checks. Always succeeds.Source§fn from_subset(element: &SS) -> SP
fn from_subset(element: &SS) -> SP
self to the equivalent element of its superset.