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
27
28
29
30
31
32
33

use num_traits::Float;

/// out[i][j] = sum_{t=j+1}^{i} x[t] for j <= i, else -infinity.
/// Returns a flattened row-major n x n matrix.
pub fn segsum<T: Float>(x: &[T]) -> Vec<T> {
	let n = x.len();
	let mut out = vec![T::neg_infinity(); n * n];
	for i in 0..n {
		out[i * n + i] = T::zero(); // empty sum on the diagonal
		let mut running = T::zero();
		for j in (0..i).rev() {
			running = running + x[j + 1];
			out[i * n + j] = running;
		}
	}
	out
}

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

	#[test]
	fn small_case() {
		let x = [1.0, 2.0, 3.0];
		let m = segsum(&x);
		// row 2, col 0 = x[1] + x[2] = 5.0
		assert_eq!(m[2 * 3 + 0], 5.0);
		// diagonal is always 0
		assert_eq!(m[0], 0.0);
		assert_eq!(m[1 * 3 + 1], 0.0);
	}
}