zilla-muf 0.1.0

Shared structured-matrix and numerical primitives for sparse attention and state space models (SSMs).
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

use num_traits::Float;

/// Sequential (recurrent) scan: h_t = a_t * h_{t-1} + b_t
pub fn sequential_scan<T: Float>(a: &[T], b: &[T], h0: T) -> Vec<T> {
	assert_eq!(a.len(), b.len(), "a and b must be the same length");
	let mut h = h0;
	let mut out = Vec::with_capacity(a.len());
	for i in 0..a.len() {
		h = a[i] * h + b[i];
		out.push(h);
	}
	out
}

#[cfg(test)]
mod tests {
	use super::*;

	#[test]
	fn cumulative_sum_case() {
		// a = 1 everywhere -> plain running sum
		let a = [1.0, 1.0, 1.0, 1.0];
		let b = [1.0, 2.0, 3.0, 4.0];
		assert_eq!(sequential_scan(&a, &b, 0.0), vec![1.0, 3.0, 6.0, 10.0]);
	}
}