oxfuse_core/
op.rs

1//! The operation set. Adding an op means adding a variant here, then letting
2//! the compiler list every exhaustive `match` that needs a new arm.
3
4#[derive(Debug, Clone, PartialEq, Eq, Hash)]
5pub enum Op {
6    Relu,
7    Sigmoid,
8    Add,
9    Mul,
10    MatMul,
11    Reshape { shape: Vec<usize> },
12    Transpose { perm: Vec<usize> },
13}
14
15impl Op {
16    /// Every op name the loader accepts, in declaration order. The unit test
17    /// below fails when a variant is added without updating this list.
18    pub const ALL_NAMES: [&'static str; 7] = [
19        "Relu",
20        "Sigmoid",
21        "Add",
22        "Mul",
23        "MatMul",
24        "Reshape",
25        "Transpose",
26    ];
27
28    pub fn name(&self) -> &'static str {
29        match self {
30            Op::Relu => "Relu",
31            Op::Sigmoid => "Sigmoid",
32            Op::Add => "Add",
33            Op::Mul => "Mul",
34            Op::MatMul => "MatMul",
35            Op::Reshape { .. } => "Reshape",
36            Op::Transpose { .. } => "Transpose",
37        }
38    }
39
40    pub fn arity(&self) -> usize {
41        match self {
42            Op::Relu | Op::Sigmoid | Op::Reshape { .. } | Op::Transpose { .. } => 1,
43            Op::Add | Op::Mul | Op::MatMul => 2,
44        }
45    }
46}
47
48#[cfg(test)]
49mod tests {
50    use super::*;
51
52    fn one_of_each() -> Vec<Op> {
53        vec![
54            Op::Relu,
55            Op::Sigmoid,
56            Op::Add,
57            Op::Mul,
58            Op::MatMul,
59            Op::Reshape { shape: vec![1] },
60            Op::Transpose { perm: vec![0] },
61        ]
62    }
63
64    /// No wildcard arm: adding a variant breaks the build here, which points
65    /// at `one_of_each` and `ALL_NAMES` above.
66    fn assert_listed(op: &Op) {
67        match op {
68            Op::Relu
69            | Op::Sigmoid
70            | Op::Add
71            | Op::Mul
72            | Op::MatMul
73            | Op::Reshape { .. }
74            | Op::Transpose { .. } => {}
75        }
76    }
77
78    #[test]
79    fn all_names_lists_every_variant_once() {
80        let ops = one_of_each();
81        ops.iter().for_each(assert_listed);
82        let names: Vec<&str> = ops.iter().map(Op::name).collect();
83        assert_eq!(names, Op::ALL_NAMES);
84    }
85}
oxfuse_core/op.rs

oxfuse_core/
op.rs
