48 lines
1.4 KiB
Rust
48 lines
1.4 KiB
Rust
//! Propagation loop in Rust — tight versor_apply chain.
|
|
//!
|
|
//! propagate_n steps runs N versor_apply calls in a single Rust stack frame,
|
|
//! eliminating Python dispatch overhead for each step.
|
|
//! Used by generate/stream.py when stepping more than one token at a time
|
|
//! (e.g. prefill, speculative steps, or batch generation).
|
|
|
|
use crate::versor::versor_apply_raw;
|
|
use thiserror::Error;
|
|
|
|
#[derive(Debug, Error)]
|
|
pub enum PropagateError {
|
|
#[error("Versor error during propagation: {0}")]
|
|
Versor(String),
|
|
}
|
|
|
|
/// Run n versor_apply steps in sequence.
|
|
/// rotors: slice of n [f32;32] versors to apply in order
|
|
/// f0: initial field state
|
|
/// Returns final field state after n steps.
|
|
pub fn propagate_n_raw(
|
|
rotors: &[[f32; 32]],
|
|
f0: &[f32; 32],
|
|
) -> Result<[f32; 32], PropagateError> {
|
|
let mut f = *f0;
|
|
for v in rotors {
|
|
f = versor_apply_raw(v, &f)
|
|
.map_err(|e| PropagateError::Versor(e.to_string()))?;
|
|
}
|
|
Ok(f)
|
|
}
|
|
|
|
/// Parallel batch propagation: apply the same rotor V to a batch of field states.
|
|
/// Used for beam search or multi-hypothesis generation.
|
|
/// Returns new batch of field states.
|
|
pub fn propagate_batch_raw(
|
|
v: &[f32; 32],
|
|
fields: &[[f32; 32]],
|
|
) -> Result<Vec<[f32; 32]>, PropagateError> {
|
|
use rayon::prelude::*;
|
|
fields
|
|
.par_iter()
|
|
.map(|f| {
|
|
versor_apply_raw(v, f)
|
|
.map_err(|e| PropagateError::Versor(e.to_string()))
|
|
})
|
|
.collect()
|
|
}
|