Struct SteTrainer 

pub struct SteTrainer {
    pub w1_latent: Vec<f32>,
    pub w2_latent: Vec<f32>,
    pub in_features: usize,
    pub hidden_size: usize,
    pub out_features: usize,
    pub config: QatConfig,
}

Maintains latent f32 shadow weights for a 2-layer MLP. During each training step:

1. Quantize latent weights → ternary {-1, 0, +1}
2. Forward pass (f32 arithmetic on quantized weights)
3. MSE loss + backprop through STE
4. SGD update on latent weights

Fields

w1_latent: Vec<f32>

w2_latent: Vec<f32>

in_features: usize

hidden_size: usize

out_features: usize

config: QatConfig

Implementations

impl SteTrainer

pub fn from_mlp(mlp: &TernaryMLP, config: QatConfig) -> Self

Initialise from an existing TernaryMLP’s quantized weights. The ternary {-1,0,+1} values become the initial latent floats.

pub fn from_f32( in_features: usize, hidden_size: usize, out_features: usize, w1_f32: Vec<f32>, w2_f32: Vec<f32>, config: QatConfig, ) -> Self

Initialise from raw f32 weights (quantization happens at first step).

pub fn train_step(&mut self, input: &[f32], target: &[f32]) -> f32

One SGD step on a single sample.

input — flat f32 row vector, length in_features target — flat f32 row vector, length out_features

Returns the MSE loss for this sample.

pub fn train(&mut self, samples: &[(Vec<f32>, Vec<f32>)]) -> QatResult

Run the full training loop over a dataset.

samples — slice of (input, target) pairs

pub fn finalize(&self) -> TernaryMLP

Finalise training: quantize latent weights and return a TernaryMLP.