#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum SaeFitLane {
DenseCertification,
SparseCodes,
CurvedStreaming,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct SaeFitAdmission {
pub lane: SaeFitLane,
pub n_obs: usize,
pub output_dim: usize,
pub n_atoms: usize,
pub dense_assignment_cells: usize,
pub response_cells: usize,
}
impl SaeFitAdmission {
pub fn uses_sparse_codes(&self) -> bool {
self.lane == SaeFitLane::SparseCodes
}
}
pub fn admit_sae_fit(
n_obs: usize,
output_dim: usize,
n_atoms: usize,
) -> Result<SaeFitAdmission, String> {
if n_obs == 0 || output_dim == 0 || n_atoms == 0 {
return Err(format!(
"admit_sae_fit requires positive N, P, and K; got N={n_obs}, P={output_dim}, K={n_atoms}"
));
}
let dense_assignment_cells = n_obs.saturating_mul(n_atoms);
let response_cells = n_obs.saturating_mul(output_dim);
let lane = if dense_assignment_cells <= response_cells {
SaeFitLane::DenseCertification
} else {
SaeFitLane::SparseCodes
};
Ok(SaeFitAdmission {
lane,
n_obs,
output_dim,
n_atoms,
dense_assignment_cells,
response_cells,
})
}
pub fn admit_topk_manifold(
n_obs: usize,
output_dim: usize,
n_atoms: usize,
d_max: usize,
support_k: usize,
) -> Result<SaeFitAdmission, String> {
let budget_bytes = crate::manifold::sae_host_in_core_budget_bytes().0;
admit_topk_manifold_with_budget(n_obs, output_dim, n_atoms, d_max, support_k, budget_bytes)
}
pub(crate) fn admit_topk_manifold_with_budget(
n_obs: usize,
output_dim: usize,
n_atoms: usize,
d_max: usize,
support_k: usize,
budget_bytes: usize,
) -> Result<SaeFitAdmission, String> {
let admission = admit_sae_fit(n_obs, output_dim, n_atoms)?;
if support_k == 0 || support_k > n_atoms {
return Err(format!(
"admit_topk_manifold requires 1 <= support_k <= K={n_atoms}; got {support_k}"
));
}
if d_max == 0 {
return Err("admit_topk_manifold requires d_max >= 1".to_string());
}
if admission.lane == SaeFitLane::DenseCertification {
return Ok(admission);
}
let ledger = crate::manifold::sae_topk_curved_budget_from_budget(
n_obs,
output_dim,
n_atoms,
d_max,
support_k,
budget_bytes,
);
if ledger.resident_seed_admitted || ledger.streaming_admitted {
return Ok(SaeFitAdmission {
lane: SaeFitLane::CurvedStreaming,
..admission
});
}
Err(format!(
"topk manifold engine refused: the dense seed ({} bytes) exceeds the host in-core \
budget ({budget_bytes} bytes) AND the streamed curved shape (peak {} bytes) exceeds \
the streaming budget ({} bytes) at N={n_obs}, K={n_atoms}, k_active={support_k}, \
d_max={d_max}. Reduce n_obs (row-subsample; the HT outer subsampling keeps the \
criterion honest) or reduce support_k — a TOPK MANIFOLD request is never silently \
substituted with the linear sparse-code lane",
ledger.resident_seed_bytes, ledger.streaming_peak_bytes, ledger.streaming_budget_bytes,
))
}
pub fn admit_linear_dictionary(
n_obs: usize,
output_dim: usize,
n_atoms: usize,
block_size: usize,
) -> Result<SaeFitAdmission, String> {
if n_obs == 0 || output_dim == 0 || n_atoms == 0 {
return Err(format!(
"admit_linear_dictionary requires positive N, P, and K; got N={n_obs}, \
P={output_dim}, K={n_atoms}"
));
}
if block_size == 0 || block_size > output_dim {
return Err(format!(
"admit_linear_dictionary requires 1 <= block_size <= P={output_dim} (a block's b \
orthonormal directions must fit in R^P); got {block_size}"
));
}
Ok(SaeFitAdmission {
lane: SaeFitLane::SparseCodes,
n_obs,
output_dim,
n_atoms,
dense_assignment_cells: n_obs.saturating_mul(n_atoms),
response_cells: n_obs.saturating_mul(output_dim),
})
}
pub fn admit_crosscoder_border(
border_dim: usize,
full_beta_dim: usize,
budget_bytes: usize,
) -> Result<(), String> {
let border_bytes = border_dim
.saturating_mul(border_dim)
.saturating_mul(std::mem::size_of::<f64>());
if border_bytes <= budget_bytes {
return Ok(());
}
let full_bytes = full_beta_dim
.saturating_mul(full_beta_dim)
.saturating_mul(std::mem::size_of::<f64>());
Err(format!(
"crosscoder border refused: the arrow-Schur border workspace at the stacked width is \
{border_dim}² · 8 = {border_bytes} bytes, over the host in-core budget \
({budget_bytes} bytes); the dense full-B border at this shape is (Σ M_k·p̃)² · 8 = \
{full_bytes} bytes. Default the atoms onto profiled Grassmann frames \
(maybe_activate_decoder_frame: the factored border Σ M_k·r_k is p̃-independent) or \
reduce the stacked width — a crosscoder request is never silently narrowed to fewer \
layers"
))
}
pub fn admit_dense_certification(
n_obs: usize,
output_dim: usize,
n_atoms: usize,
) -> Result<SaeFitAdmission, String> {
let admission = admit_sae_fit(n_obs, output_dim, n_atoms)?;
if admission.uses_sparse_codes() {
return Err(format!(
"dense manifold engine refused: it is the small-K certification lane \
(admitted only while N*K <= N*P, i.e. K <= P); N={n_obs}, P={output_dim}, K={n_atoms} \
gives N*K={} > N*P={} — the overcomplete (K > P) routes are the hard TOP-K SUPPORT \
curved lane (assignment='topk', admitted by concrete memory budget; penalty-gated \
modes cannot take it because their N×K gate logits are live Newton state) or the \
linear sparse-code lane; neither builds the dense N×K assignment state",
admission.dense_assignment_cells, admission.response_cells
));
}
Ok(admission)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn topk_manifold_admits_overcomplete_to_curved_lane_and_refuses_over() {
let bytes = 2 * 4096usize * 32_000 * 8;
let ok = admit_topk_manifold_with_budget(4096, 512, 32_000, 1, 8, bytes)
.expect("within-budget topk admission");
assert_eq!(
ok.lane,
SaeFitLane::CurvedStreaming,
"TopK within budget admits the CURVED manifold lane at K > P"
);
assert_eq!((ok.n_obs, ok.output_dim, ok.n_atoms), (4096, 512, 32_000));
let ledger =
crate::manifold::sae_topk_curved_budget_from_budget(4096, 512, 32_000, 1, 8, bytes);
assert!(ledger.resident_seed_admitted);
assert_eq!(ledger.resident_seed_bytes, bytes);
let err = admit_topk_manifold_with_budget(4096, 512, 32_000, 1, 8, bytes - 1)
.expect_err("over-both-budgets topk must refuse");
assert!(
err.contains("never silently substituted"),
"refusal must state the no-substitution contract; got: {err}"
);
let softmax = admit_sae_fit(4096, 512, 32_000).expect("softmax admission");
assert_eq!(softmax.lane, SaeFitLane::SparseCodes);
let small = admit_topk_manifold_with_budget(1024, 4096, 128, 2, 4, 0)
.expect("K <= P is admitted regardless of budget");
assert_eq!(small.lane, SaeFitLane::DenseCertification);
assert_eq!(
small,
admit_sae_fit(1024, 4096, 128).expect("dense admission")
);
assert!(admit_topk_manifold_with_budget(64, 8, 16, 1, 0, bytes).is_err());
assert!(admit_topk_manifold_with_budget(64, 8, 16, 1, 17, bytes).is_err());
}
#[test]
fn topk_manifold_streaming_region_admits_to_curved_lane() {
let ledger = crate::manifold::sae_topk_curved_budget_from_budget(
1_000_000,
512,
32_000,
1,
8,
16 * 1024 * 1024 * 1024,
);
assert!(!ledger.resident_seed_admitted);
assert!(ledger.streaming_admitted);
let admission =
admit_topk_manifold_with_budget(1_000_000, 512, 32_000, 1, 8, 16 * 1024 * 1024 * 1024)
.expect("streaming region is runnable and must admit to the curved lane");
assert_eq!(
admission.lane,
SaeFitLane::CurvedStreaming,
"the over-resident/under-streaming region admits to CurvedStreaming"
);
assert_eq!(
(admission.n_obs, admission.output_dim, admission.n_atoms),
(1_000_000, 512, 32_000)
);
let err = admit_topk_manifold_with_budget(1_000_000, 512, 32_000, 1, 8, 0)
.expect_err("over-both-budgets topk must refuse");
assert!(err.contains("never silently substituted"));
assert!(!err.contains("admit_topk_curved_lane"));
}
#[test]
fn explicit_linear_dictionary_admits_sparse_codes_at_any_k() {
let small = admit_linear_dictionary(640, 24, 16, 1).expect("K<=P linear admission");
assert_eq!(small.lane, SaeFitLane::SparseCodes);
assert_eq!(
(small.n_obs, small.output_dim, small.n_atoms),
(640, 24, 16)
);
let large = admit_linear_dictionary(4096, 512, 32_000, 1).expect("K>P linear admission");
assert_eq!(large.lane, SaeFitLane::SparseCodes);
let block = admit_linear_dictionary(4096, 512, 1024, 4).expect("block linear admission");
assert_eq!(block.lane, SaeFitLane::SparseCodes);
assert!(admit_linear_dictionary(4096, 512, 1024, 513).is_err());
assert!(admit_linear_dictionary(4096, 512, 1024, 0).is_err());
assert!(admit_linear_dictionary(0, 512, 1024, 1).is_err());
}
#[test]
fn crosscoder_border_admission_checks_actual_border_and_names_frames() {
admit_crosscoder_border(100, 100, 100 * 100 * 8).expect("within-budget border");
admit_crosscoder_border(512, 4096, 512 * 512 * 8).expect("framed border within budget");
let err = admit_crosscoder_border(4096, 4096, 512 * 512 * 8)
.expect_err("full-B border over budget must refuse");
assert!(
err.contains("never silently narrowed"),
"refusal must state the no-narrowing contract; got: {err}"
);
assert!(
err.contains("maybe_activate_decoder_frame"),
"refusal must name the frame remedy; got: {err}"
);
}
#[test]
fn admission_demotes_dense_when_assignment_state_exceeds_response() {
let small = admit_sae_fit(1024, 4096, 128).expect("small admission");
assert_eq!(small.lane, SaeFitLane::DenseCertification);
let large = admit_sae_fit(1024, 4096, 32_000).expect("large admission");
assert_eq!(large.lane, SaeFitLane::SparseCodes);
assert!(large.dense_assignment_cells > large.response_cells);
}
#[test]
fn dense_certification_admits_k_le_p_and_refuses_sparse_k_gt_p() {
let ok = admit_dense_certification(1024, 4096, 128).expect("K<P must be admitted");
assert_eq!(ok.lane, SaeFitLane::DenseCertification);
assert!(
admit_dense_certification(64, 8, 8).is_ok(),
"K == P is the DenseCertification boundary and must be admitted"
);
let err = admit_dense_certification(1_000_000, 4096, 32_000)
.expect_err("K > P must be refused by the dense-engine guard");
assert!(
err.contains("sparse-code lane"),
"the refusal must point at the sparse-code lane; got: {err}"
);
assert!(admit_dense_certification(10, 100, 100).is_ok());
assert!(admit_dense_certification(10, 100, 101).is_err());
assert!(admit_dense_certification(0, 4, 3).is_err());
}
}