tokitai-operator 0.1.0

Verified DL kernel compiler: formally-checked GEMM, p-adic, sheaf, contract-carrying ops. Paper-artifact grade.
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
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100

use tokitai_operator::backend::TensorStore;
use tokitai_operator::facade::Tokitai;
use tokitai_operator::object::{ObjectMeta, Representation, Shape, Tensor};
use tokitai_operator::op::arithmetic::SquareOp;
use tokitai_operator::planner::PlanStepKind;

fn integer_meta(shape: Vec<usize>) -> ObjectMeta {
    ObjectMeta::tensor(
        tokitai_operator::domain::DomainId::new("integer"),
        Shape::from(shape),
        Representation::dense_cpu(),
    )
}

#[test]
fn square_planner_lowers_single_input_to_single_step() {
    let api = Tokitai::cpu_only();
    let mut builder = tokitai_operator::facade::FacadeGraphBuilder::new();
    let data = builder.input(integer_meta(vec![5]));
    let out = builder.add_op(SquareOp, &[data]).expect("square infer")[0];
    let graph = builder.build();
    let plan = api.plan_public(&graph).expect("plan should succeed");

    let single_steps: Vec<_> = plan
        .plan
        .steps
        .iter()
        .filter(|step| step.kind == PlanStepKind::Single && step.op_name == "square")
        .collect();
    assert_eq!(single_steps.len(), 1);
    assert_eq!(
        single_steps[0].lowering_rule_id.as_deref(),
        Some("cpu.square.dense")
    );

    let mut store = TensorStore::<i64>::new();
    let int_domain = tokitai_operator::domain::DomainId::new("integer");
    store.insert(
        data,
        Tensor::dense_cpu(int_domain, Shape::from(vec![5]), vec![-3, -1, 0, 1, 3]),
    );
    let (executed, _report) = api
        .execute_i64(&graph, &plan.plan, store, &[out])
        .expect("square execution should succeed");
    let result = executed.get(out).unwrap();
    assert_eq!(result.data, vec![9, 1, 0, 1, 9]);
}

#[test]
fn square_explicit_per_element_reference() {
    let api = Tokitai::cpu_only();
    let mut builder = tokitai_operator::facade::FacadeGraphBuilder::new();
    let data = builder.input(integer_meta(vec![6]));
    let out = builder.add_op(api.square(), &[data]).unwrap()[0];
    let graph = builder.build();
    let plan = api.plan_public(&graph).unwrap();

    let input_data = vec![-7, -3, 0, 4, 5, 12];
    let mut store = TensorStore::<i64>::new();
    let int_domain = tokitai_operator::domain::DomainId::new("integer");
    store.insert(
        data,
        Tensor::dense_cpu(int_domain, Shape::from(vec![6]), input_data.clone()),
    );
    let (executed, _report) = api.execute_i64(&graph, &plan.plan, store, &[out]).unwrap();
    let result = executed.get(out).unwrap().data.clone();
    let expected: Vec<i64> = input_data.iter().map(|v| v.wrapping_mul(*v)).collect();
    assert_eq!(result, expected);
}

#[test]
fn square_with_empty_data_returns_empty() {
    let api = Tokitai::cpu_only();
    let mut builder = tokitai_operator::facade::FacadeGraphBuilder::new();
    let data = builder.input(integer_meta(vec![0]));
    let out = builder.add_op(api.square(), &[data]).unwrap()[0];
    let graph = builder.build();
    let plan = api.plan_public(&graph).unwrap();
    let mut store = TensorStore::<i64>::new();
    let int_domain = tokitai_operator::domain::DomainId::new("integer");
    store.insert(
        data,
        Tensor::dense_cpu(int_domain, Shape::from(vec![0]), vec![]),
    );
    let (executed, _report) = api.execute_i64(&graph, &plan.plan, store, &[out]).unwrap();
    let result = executed.get(out).unwrap();
    assert!(result.is_empty());
}

#[test]
fn square_planner_rejects_wrong_input_count() {
    let mut builder = tokitai_operator::facade::FacadeGraphBuilder::new();
    let a = builder.input(integer_meta(vec![2]));
    let b = builder.input(integer_meta(vec![2]));
    let err = builder
        .add_op(SquareOp, &[a, b])
        .expect_err("square should reject 2 inputs");
    let msg = format!("{err}");
    assert!(msg.contains("1 input"), "got: {msg}");
}