Skip to main content

Module row_jet_program

Module row_jet_program 

Source
Expand description

The SAE reconstruction row as a single Taylor-jet program (issue #932).

§The row program

The exact-LAML SAE engine needs, per row, the derivative tower of the reconstruction

  ẑ_row,c(p) = Σ_k ζ_k(ℓ) · decoded_{k,c}(t_k),   decoded_{k,c}(t) = Σ_b Φ_b(t)·B_{b,c}

— a gate nonlinearity ζ(ℓ) (softmax / IBP sigmoid) composed with a basis Φ(t) composed with a linear decoder B, in the per-row primary coordinates p = (gate logits ℓ, latent coordinates t). Today the arrow-Schur assembly (SaeManifoldTerm::row_jets_for_logdet) hand-packs the first/second channels of this reconstruction from separate gate derivative arrays (gate_derivatives_for_row) and basis jet tensors — exactly the kind of hand-maintained cross-block tower whose sign flips are the #736 / desync bug genus. The #1006 third-order logdet adjoint Γ_a = tr(H⁻¹ ∂H/∂θ_a) is the consumer of those very channels.

This module writes that reconstruction once over the Tower4<K> scalar so the value/gradient/Hessian/third channels of one row come from ONE jet evaluation. SaeReconstructionRowProgram is generic over the gate kind and the per-row basis jets; the gate, basis and decoder compose with plain Tower4 arithmetic, so there is no separate “channel” to forget.

§The basis as a local jet

The production assembly does NOT re-evaluate the manifold basis Φ as a function of perturbed coordinates: it consumes the precomputed jet tensors (Φ, ∂Φ/∂t, ∂²Φ/∂t²) evaluated at the current t. The reconstruction’s dependence on t is therefore defined by those tensors — the local quadratic Taylor model of Φ about the current point. This program builds each basis function as exactly that Tower4 quadratic from the stored jets, so the value/first/second channels it emits are the same object the hand path packs — derived by independent arithmetic (tower Leibniz / Faà di Bruno vs hand-summed cross terms). Agreement across both is a true correctness proof of the hand kernel; disagreement names a dropped or sign-flipped cross block loudly. That oracle is the riding test below.

Structs§

AtomRowBasisJet
One atom’s local basis jet at the current row: the stored (value, jacobian, second) jet tensors of Φ plus the decoder block B. Indexed [basis_col], [basis_col][axis], [basis_col][axis_a][axis_b], and [basis_col][out_col].
SaeReconstructionRowProgram
One row of the SAE reconstruction as a jet program: the per-atom basis jets, the gate, the current gate-logit values, and the primary layout that maps (atom logit, atom latent axis) to a seeded tower variable slot.

Enums§

RowGate
The gate nonlinearity ζ(ℓ) of the SAE assignment, as the row program sees it. The production term carries the same two smooth branches (softmax over a shared partition; per-atom IBP/JumpReLU sigmoid); the program reproduces the branch the criterion evaluates so the value channel is the production gate.

Constants§

SAE_FIXED_COORD_SLOT
Sentinel in SaeReconstructionRowProgram::coord_slot for an atom coordinate that is fixed in this row’s local chart (compact active-set rows omit inactive atom coordinates, but softmax logit derivatives can still see that atom’s decoded value as a constant).