oxionnx 0.1.2

Pure Rust ONNX inference engine — zero C/C++ dependencies
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
101
102
103
104
105

//! Operator classification for mixed precision inference.
//!
//! Determines whether an operator can run in f16 or requires f32 for numerical stability.

/// Returns `true` if the operator can safely execute with f16 inputs/outputs
/// without significant precision loss.
pub fn should_use_f16(op_type: &str) -> bool {
    matches!(
        op_type,
        // Element-wise activations & arithmetic
        "Add"
            | "Sub"
            | "Mul"
            | "Div"
            | "Relu"
            | "LeakyRelu"
            | "Sigmoid"
            | "Tanh"
            | "Gelu"
            | "Silu"
            | "SiLU"
            | "HardSigmoid"
            | "HardSwish"
            | "Abs"
            | "Neg"
            | "Sqrt"
            | "Reciprocal"
            | "Clip"
            | "Erf"
            | "Softsign"
            | "Softplus"
            | "Mish"
            | "Celu"
            | "Elu"
            | "Selu"
            | "ThresholdedRelu"
            | "PRelu"
            // Normalization (compute in f16; gamma/beta stay f32 but output is f16)
            | "LayerNormalization"
            | "LayerNorm"
            | "BatchNormalization"
            | "BatchNorm"
            | "GroupNormalization"
            | "GroupNorm"
            | "RMSNorm"
            | "SimplifiedLayerNormalization"
            | "InstanceNorm"
            | "InstanceNormalization"
            // Softmax (stable computation works in f16)
            | "Softmax"
            | "LogSoftmax"
            // Shape manipulation (zero-copy or simple data movement)
            | "Transpose"
            | "Reshape"
            | "Concat"
            | "Slice"
            | "Split"
            | "Squeeze"
            | "Unsqueeze"
            | "Flatten"
            | "Identity"
            | "Expand"
            | "Tile"
            | "DepthToSpace"
            | "SpaceToDepth"
            // Attention (scores in f16)
            | "Attention"
            | "MultiHeadAttention"
            | "RotaryEmbedding"
            // Dropout (just passthrough or mask)
            | "Dropout"
    )
}

/// Returns `true` if the operator requires f32 for numerical stability
/// (accumulation, precision-sensitive math).
pub fn requires_f32(op_type: &str) -> bool {
    matches!(
        op_type,
        "MatMul"
            | "Gemm"
            | "ReduceSum"
            | "ReduceMean"
            | "ReduceMax"
            | "ReduceMin"
            | "ReduceProd"
            | "ReduceL1"
            | "ReduceL2"
            | "ReduceLogSum"
            | "ReduceLogSumExp"
            | "ReduceSumSquare"
            | "Pow"
            | "Exp"
            | "Log"
            | "Conv"
            | "ConvTranspose"
            | "ConvAddRelu"
            | "MaxPool"
            | "AveragePool"
            | "GlobalAveragePool"
            | "GlobalMaxPool"
            | "CumSum"
            | "Einsum"
    )
}