knok 0.1.0

Compile-time linalg graphs for Rust
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
34

use knok::prelude::*;
use knok::{Engine, RuntimeConfig};

#[knok::graph(backend = "llvm-cpu")]
fn classifier_head(logits: Tensor1<f32, 4>) -> Tensor1<f32, 4> {
    softmax(logits)
}

#[knok::graph(backend = "llvm-cpu")]
fn predicted_class(logits: Tensor1<f32, 4>) -> Tensor1<f32, 1> {
    argmax(logits)
}

fn main() -> knok::Result<()> {
    let engine = Engine::new(RuntimeConfig::auto())?;
    let logits = Tensor1::from_array([1.0, 2.0, 3.0, 4.0]);

    let probabilities = classifier_head_run(&engine, logits.clone())?;
    let class = predicted_class_run(&engine, logits)?;

    assert_close(
        &probabilities.into_vec(),
        &[0.032058604, 0.08714432, 0.23688284, 0.6439143],
    );
    assert_eq!(class.into_vec(), vec![3.0]);
    Ok(())
}

fn assert_close(actual: &[f32], expected: &[f32]) {
    assert_eq!(actual.len(), expected.len());
    for (actual, expected) in actual.iter().zip(expected) {
        assert!((actual - expected).abs() < 1.0e-4);
    }
}