knok-build 0.4.0

Build-script graph tracing and code generation for knok
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
101
102

use knok_core::CallOp;

use super::{
    expr::call_output,
    tensor::{TraceOperand, TraceTensor},
};

/// Records a static `arange(0, stop, 1)` with shape and dtype from `Output`.
pub fn arange_to<Output>(stop: impl TraceOperand<Output::Elem>) -> Output
where
    Output: TraceTensor,
{
    trace_arange1(stop)
}

fn trace_arange1<Output>(stop: impl TraceOperand<Output::Elem>) -> Output
where
    Output: TraceTensor,
{
    call_output(
        CallOp::Arange(Output::tensor_type()),
        vec![stop.into_operand_expr()],
    )
}

/// Records a static `arange(start, stop, 1)` with shape and dtype from `Output`.
pub fn arange<Output>(
    start: impl TraceOperand<Output::Elem>,
    stop: impl TraceOperand<Output::Elem>,
) -> Output
where
    Output: TraceTensor,
{
    trace_arange2(start, stop)
}

fn trace_arange2<Output>(
    start: impl TraceOperand<Output::Elem>,
    stop: impl TraceOperand<Output::Elem>,
) -> Output
where
    Output: TraceTensor,
{
    call_output(
        CallOp::Arange(Output::tensor_type()),
        vec![start.into_operand_expr(), stop.into_operand_expr()],
    )
}

/// Records a static `arange(start, stop, step)` with shape and dtype from `Output`.
pub fn arange_step<Output>(
    start: impl TraceOperand<Output::Elem>,
    stop: impl TraceOperand<Output::Elem>,
    step: impl TraceOperand<Output::Elem>,
) -> Output
where
    Output: TraceTensor,
{
    trace_arange3(start, stop, step)
}

fn trace_arange3<Output>(
    start: impl TraceOperand<Output::Elem>,
    stop: impl TraceOperand<Output::Elem>,
    step: impl TraceOperand<Output::Elem>,
) -> Output
where
    Output: TraceTensor,
{
    call_output(
        CallOp::Arange(Output::tensor_type()),
        vec![
            start.into_operand_expr(),
            stop.into_operand_expr(),
            step.into_operand_expr(),
        ],
    )
}

/// Records a static `linspace(start, stop)` with shape and dtype from `Output`.
pub fn linspace<Output>(
    start: impl TraceOperand<Output::Elem>,
    stop: impl TraceOperand<Output::Elem>,
) -> Output
where
    Output: TraceTensor,
{
    call_output(
        CallOp::Linspace(Output::tensor_type()),
        vec![start.into_operand_expr(), stop.into_operand_expr()],
    )
}

/// Records a static identity matrix with shape and dtype from `Output`.
pub fn eye<Output: TraceTensor>() -> Output {
    call_output(CallOp::Eye(Output::tensor_type()), Vec::new())
}

/// Alias for [`eye`].
pub fn identity<Output: TraceTensor>() -> Output {
    eye()
}