teeny_onnx/
onnx.rs

1/*
2 * Copyright (c) 2026 Teenygrad.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *   http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17use std::{
18    collections::{BTreeMap, BTreeSet},
19    fs::File,
20    io::Read,
21    path::Path,
22};
23
24use protobuf::{CodedInputStream, Enum, Message};
25use teeny_core::graph::{DtypeRepr, Graph, Op, Shape};
26
27use crate::{
28    errors::{Error, Result},
29    onnx::onnx_proto3::{ModelProto, NodeProto, TensorProto, ValueInfoProto, tensor_proto},
30};
31
32include!(concat!(env!("OUT_DIR"), "/protos/mod.rs"));
33
34/// Parses ONNX protobuf models into a `teeny-core` [`Graph`].
35pub struct Onnx {}
36
37impl Onnx {
38    /// Reads and parses the `.onnx` file at `path`.
39    pub fn from_path(path: impl AsRef<Path>) -> Result<Graph> {
40        let mut file = File::open(path)?;
41        Self::from_reader(&mut file)
42    }
43
44    /// Parses an ONNX model (`ModelProto`, protobuf binary format) from `reader` into a `Graph`.
45    ///
46    /// Errors if the bytes aren't valid ONNX protobuf, the model has no graph, or the graph uses
47    /// an op this crate doesn't support yet.
48    pub fn from_reader(reader: &mut impl Read) -> Result<Graph> {
49        let mut reader = CodedInputStream::new(reader);
50        let model = ModelProto::parse_from(&mut reader)?;
51        let onnx_graph = model
52            .graph
53            .as_ref()
54            .ok_or_else(|| Error::InvalidModel("Model is missing graph field".to_string()))?;
55
56        let mut graph = Graph::new();
57        let mut value_to_node: BTreeMap<String, usize> = BTreeMap::new();
58
59        let mut value_info_map: BTreeMap<String, (DtypeRepr, Shape)> = BTreeMap::new();
60        for info in &onnx_graph.input {
61            if let Some(meta) = value_info_meta(info)? {
62                value_info_map.insert(info.name.clone(), meta);
63            }
64        }
65        for info in &onnx_graph.value_info {
66            if let Some(meta) = value_info_meta(info)? {
67                value_info_map.insert(info.name.clone(), meta);
68            }
69        }
70        for info in &onnx_graph.output {
71            if let Some(meta) = value_info_meta(info)? {
72                value_info_map.insert(info.name.clone(), meta);
73            }
74        }
75
76        let initializer_names: BTreeSet<&str> = onnx_graph
77            .initializer
78            .iter()
79            .map(|t| t.name.as_str())
80            .collect();
81        let initializer_map: BTreeMap<&str, &TensorProto> = onnx_graph
82            .initializer
83            .iter()
84            .map(|t| (t.name.as_str(), t))
85            .collect();
86
87        for init in &onnx_graph.initializer {
88            let dtype = tensor_dtype_to_dtype(init.data_type)?;
89            let shape = dims_to_shape(&init.dims);
90            let id = graph.add_node(Op::Input, vec![], dtype, shape);
91            value_to_node.insert(init.name.clone(), id);
92        }
93
94        for input in &onnx_graph.input {
95            if initializer_names.contains(input.name.as_str()) {
96                continue;
97            }
98            let (dtype, shape) = value_info_map
99                .get(&input.name)
100                .cloned()
101                .unwrap_or((DtypeRepr::F32, vec![]));
102            let id = graph.add_node(Op::Input, vec![], dtype, shape);
103            value_to_node.insert(input.name.clone(), id);
104        }
105
106        for node in &onnx_graph.node {
107            let mut inputs = Vec::with_capacity(node.input.len());
108            for input_name in &node.input {
109                if input_name.is_empty() {
110                    continue;
111                }
112                let input_id = value_to_node.get(input_name).copied().ok_or_else(|| {
113                    Error::InvalidModel(format!(
114                        "Node '{}' references unknown input '{}'",
115                        node.name, input_name
116                    ))
117                })?;
118                inputs.push(input_id);
119            }
120
121            // Some ops (Constant, SequenceEmpty, OptionalHasElement) have zero
122            // tensor inputs by design — allow them through.
123            let zero_input_allowed = matches!(
124                node.op_type.as_str(),
125                "Constant" | "SequenceEmpty" | "OptionalHasElement"
126            );
127            if inputs.is_empty() && !zero_input_allowed {
128                return Err(Error::InvalidModel(format!(
129                    "Node '{}' (op='{}') has no resolvable inputs",
130                    node.name, node.op_type
131                ))
132                .into());
133            }
134
135            let first_output = node
136                .output
137                .iter()
138                .find(|name| !name.is_empty())
139                .ok_or_else(|| {
140                    Error::InvalidModel(format!("Node '{}' has no output names", node.name))
141                })?;
142
143            let (dtype, shape) = value_info_map
144                .get(first_output)
145                .cloned()
146                .unwrap_or_else(|| {
147                    if !inputs.is_empty() {
148                        let input_node = &graph.nodes[inputs[0]];
149                        (input_node.dtype, input_node.shape.clone())
150                    } else {
151                        // Zero-input op with no value_info — fall back to F32 scalar.
152                        (DtypeRepr::F32, vec![])
153                    }
154                });
155
156            let op = map_node_op(node, &initializer_map)?;
157            let node_id = graph.add_node(op, inputs, dtype, shape);
158
159            if !node.name.is_empty() {
160                graph.names.insert(node_id, node.name.clone());
161            }
162
163            for output_name in &node.output {
164                if !output_name.is_empty() {
165                    value_to_node.insert(output_name.clone(), node_id);
166                }
167            }
168        }
169
170        Ok(graph)
171    }
172}
173
174fn value_info_meta(info: &ValueInfoProto) -> Result<Option<(DtypeRepr, Shape)>> {
175    let ty = info
176        .type_
177        .as_ref()
178        .ok_or_else(|| Error::InvalidModel("ValueInfoProto has no type".to_string()))?;
179    if !ty.has_tensor_type() {
180        return Ok(None);
181    }
182    let tensor = ty.tensor_type();
183    let dtype = tensor_dtype_to_dtype(tensor.elem_type)?;
184    let shape = tensor.shape.as_ref().map_or_else(Vec::new, |shape| {
185        shape
186            .dim
187            .iter()
188            .map(|d| {
189                if d.has_dim_value() && d.dim_value() >= 0 {
190                    Some(d.dim_value() as usize)
191                } else {
192                    None
193                }
194            })
195            .collect()
196    });
197
198    Ok(Some((dtype, shape)))
199}
200
201fn dims_to_shape(dims: &[i64]) -> Shape {
202    dims.iter()
203        .map(|d| if *d >= 0 { Some(*d as usize) } else { None })
204        .collect()
205}
206
207fn tensor_dtype_to_dtype(dtype: i32) -> Result<DtypeRepr> {
208    Ok(match tensor_proto::DataType::from_i32(dtype) {
209        Some(tensor_proto::DataType::BOOL) => DtypeRepr::Bool,
210        Some(tensor_proto::DataType::INT8) => DtypeRepr::I8,
211        Some(tensor_proto::DataType::INT16) => DtypeRepr::I16,
212        Some(tensor_proto::DataType::INT32) => DtypeRepr::I32,
213        Some(tensor_proto::DataType::INT64) => DtypeRepr::I64,
214        Some(tensor_proto::DataType::UINT8) => DtypeRepr::U8,
215        Some(tensor_proto::DataType::UINT16) => DtypeRepr::U16,
216        Some(tensor_proto::DataType::UINT32) => DtypeRepr::U32,
217        Some(tensor_proto::DataType::UINT64) => DtypeRepr::U64,
218        Some(tensor_proto::DataType::FLOAT16) => DtypeRepr::F16,
219        Some(tensor_proto::DataType::BFLOAT16) => DtypeRepr::BF16,
220        Some(tensor_proto::DataType::DOUBLE) => DtypeRepr::F64,
221        Some(tensor_proto::DataType::FLOAT) | None => DtypeRepr::F32,
222        // STRING tensors — represent as U8 (byte storage) for graph purposes.
223        Some(tensor_proto::DataType::STRING) => DtypeRepr::U8,
224        // Sub-byte and float8 types — map to nearest supported type.
225        Some(tensor_proto::DataType::FLOAT8E4M3FN)
226        | Some(tensor_proto::DataType::FLOAT8E4M3FNUZ)
227        | Some(tensor_proto::DataType::FLOAT8E5M2)
228        | Some(tensor_proto::DataType::FLOAT8E5M2FNUZ) => DtypeRepr::F16,
229        Some(_) => {
230            // Unknown or newer sub-byte types (INT4, UINT4, INT2, UINT2, FLOAT4…).
231            // Map to I8/U8 as a safe placeholder for graph-building purposes.
232            if dtype % 2 == 0 {
233                DtypeRepr::U8
234            } else {
235                DtypeRepr::I8
236            }
237        }
238    })
239}
240
241fn get_attr_int(node: &NodeProto, name: &str, default: i64) -> i64 {
242    node.attribute
243        .iter()
244        .find(|a| a.name == name)
245        .map(|a| a.i)
246        .unwrap_or(default)
247}
248
249fn get_attr_float(node: &NodeProto, name: &str, default: f32) -> f32 {
250    node.attribute
251        .iter()
252        .find(|a| a.name == name)
253        .map(|a| a.f)
254        .unwrap_or(default)
255}
256
257fn get_attr_ints<'a>(node: &'a NodeProto, name: &str) -> Option<&'a [i64]> {
258    node.attribute
259        .iter()
260        .find(|a| a.name == name)
261        .map(|a| a.ints.as_slice())
262}
263
264fn get_attr_string<'a>(node: &'a NodeProto, name: &str, default: &'a str) -> &'a str {
265    node.attribute
266        .iter()
267        .find(|a| a.name == name)
268        .and_then(|a| std::str::from_utf8(&a.s).ok())
269        .unwrap_or(default)
270}
271
272fn map_node_op(node: &NodeProto, initializers: &BTreeMap<&str, &TensorProto>) -> Result<Op> {
273    let op = match node.op_type.as_str() {
274        // ---------------------------------------------------------------
275        // Existing ops (previously supported)
276        // ---------------------------------------------------------------
277        "Relu" => Op::Relu,
278        "Sigmoid" => Op::Sigmoid,
279        "Tanh" => Op::Tanh,
280        "Add" => Op::Add,
281        "Flatten" => Op::Flatten,
282        "Softmax" => {
283            let axis = get_attr_int(node, "axis", 1);
284            Op::Softmax {
285                dim: axis.max(0) as usize,
286            }
287        }
288
289        "Conv" => {
290            let weight_name = node.input.get(1).map(|s| s.as_str()).unwrap_or("");
291            // Weight dims: [out_ch, in_ch/groups, k...] if from an initializer,
292            // otherwise inferred from the kernel_shape attribute.
293            let weight_dims: Vec<i64> = if let Some(w) = initializers.get(weight_name) {
294                w.dims.clone()
295            } else {
296                let ks = get_attr_ints(node, "kernel_shape").unwrap_or(&[1, 1]);
297                let mut d = vec![0i64, 0i64]; // out/in channels unknown
298                d.extend_from_slice(ks);
299                d
300            };
301
302            let has_bias = node.input.get(2).is_some_and(|s| !s.is_empty());
303            let strides = get_attr_ints(node, "strides").unwrap_or(&[1, 1]);
304            let pads = get_attr_ints(node, "pads").unwrap_or(&[0, 0, 0, 0]);
305            let groups = get_attr_int(node, "group", 1).max(1) as usize;
306            let weight = &weight_dims[..];
307            let ndim = weight.len();
308
309            if ndim == 3 {
310                Op::Conv1d {
311                    in_channels: weight.get(1).copied().unwrap_or(0) as usize * groups,
312                    out_channels: weight.first().copied().unwrap_or(0) as usize,
313                    kernel_l: weight.get(2).copied().unwrap_or(1) as usize,
314                    stride: strides.first().copied().unwrap_or(1).max(1) as usize,
315                    padding: pads.first().copied().unwrap_or(0).max(0) as usize,
316                    has_bias,
317                }
318            } else if ndim == 4 {
319                Op::Conv2d {
320                    in_channels: weight.get(1).copied().unwrap_or(0) as usize * groups,
321                    out_channels: weight.first().copied().unwrap_or(0) as usize,
322                    kernel_h: weight.get(2).copied().unwrap_or(1) as usize,
323                    kernel_w: weight.get(3).copied().unwrap_or(1) as usize,
324                    stride_h: strides.first().copied().unwrap_or(1).max(1) as usize,
325                    stride_w: strides.get(1).copied().unwrap_or(1).max(1) as usize,
326                    padding_h: pads.first().copied().unwrap_or(0).max(0) as usize,
327                    padding_w: pads.get(1).copied().unwrap_or(0).max(0) as usize,
328                    groups,
329                    has_bias,
330                }
331            } else if ndim == 5 {
332                Op::Conv3d {
333                    in_channels: weight.get(1).copied().unwrap_or(0) as usize * groups,
334                    out_channels: weight.first().copied().unwrap_or(0) as usize,
335                    kernel_d: weight.get(2).copied().unwrap_or(1) as usize,
336                    kernel_h: weight.get(3).copied().unwrap_or(1) as usize,
337                    kernel_w: weight.get(4).copied().unwrap_or(1) as usize,
338                    stride_d: strides.first().copied().unwrap_or(1).max(1) as usize,
339                    stride_h: strides.get(1).copied().unwrap_or(1).max(1) as usize,
340                    stride_w: strides.get(2).copied().unwrap_or(1).max(1) as usize,
341                    padding_d: pads.first().copied().unwrap_or(0).max(0) as usize,
342                    padding_h: pads.get(1).copied().unwrap_or(0).max(0) as usize,
343                    padding_w: pads.get(2).copied().unwrap_or(0).max(0) as usize,
344                    has_bias,
345                }
346            } else {
347                return Err(Error::InvalidModel(format!(
348                    "Conv node '{}': unsupported weight rank {}",
349                    node.name, ndim
350                ))
351                .into());
352            }
353        }
354
355        "ConvTranspose" => {
356            let weight_name = node.input.get(1).map(|s| s.as_str()).unwrap_or("");
357            let (out_channels, kernel_h, kernel_w) = if let Some(w) = initializers.get(weight_name)
358            {
359                let nd = w.dims.len();
360                if nd >= 4 {
361                    (w.dims[1] as usize, w.dims[2] as usize, w.dims[3] as usize)
362                } else {
363                    (0, 1, 1)
364                }
365            } else {
366                (0, 1, 1)
367            };
368            let has_bias = node.input.get(2).is_some_and(|s| !s.is_empty());
369            let strides = get_attr_ints(node, "strides").unwrap_or(&[1, 1]);
370            let pads = get_attr_ints(node, "pads").unwrap_or(&[0, 0, 0, 0]);
371            let out_pads = get_attr_ints(node, "output_padding").unwrap_or(&[0, 0]);
372            let groups = get_attr_int(node, "group", 1).max(1) as usize;
373            let in_channels = if let Some(w) = initializers.get(weight_name) {
374                w.dims.first().copied().unwrap_or(0) as usize * groups
375            } else {
376                0
377            };
378            Op::ConvTranspose {
379                in_channels,
380                out_channels,
381                kernel_h,
382                kernel_w,
383                stride_h: strides.first().copied().unwrap_or(1).max(1) as usize,
384                stride_w: strides.get(1).copied().unwrap_or(1).max(1) as usize,
385                padding_h: pads.first().copied().unwrap_or(0).max(0) as usize,
386                padding_w: pads.get(1).copied().unwrap_or(0).max(0) as usize,
387                output_padding_h: out_pads.first().copied().unwrap_or(0).max(0) as usize,
388                output_padding_w: out_pads.get(1).copied().unwrap_or(0).max(0) as usize,
389                groups,
390                has_bias,
391            }
392        }
393
394        "MaxPool" => {
395            let kernel = get_attr_ints(node, "kernel_shape").unwrap_or(&[1, 1]);
396            let strides = get_attr_ints(node, "strides").unwrap_or(kernel);
397            let pads = get_attr_ints(node, "pads").unwrap_or(&[0, 0, 0, 0]);
398            match kernel.len() {
399                1 => Op::MaxPool1d {
400                    kernel_l: kernel[0].max(1) as usize,
401                    stride: strides.first().copied().unwrap_or(kernel[0]).max(1) as usize,
402                },
403                3 => Op::MaxPool3d {
404                    kernel_d: kernel[0].max(1) as usize,
405                    kernel_h: kernel[1].max(1) as usize,
406                    kernel_w: kernel[2].max(1) as usize,
407                    stride_d: strides.first().copied().unwrap_or(kernel[0]).max(1) as usize,
408                    stride_h: strides.get(1).copied().unwrap_or(kernel[1]).max(1) as usize,
409                    stride_w: strides.get(2).copied().unwrap_or(kernel[2]).max(1) as usize,
410                },
411                _ => Op::MaxPool2d {
412                    kernel_h: kernel.first().copied().unwrap_or(1).max(1) as usize,
413                    kernel_w: kernel.get(1).copied().unwrap_or(1).max(1) as usize,
414                    stride_h: strides.first().copied().unwrap_or(1).max(1) as usize,
415                    stride_w: strides.get(1).copied().unwrap_or(1).max(1) as usize,
416                    pad_h: pads.first().copied().unwrap_or(0).max(0) as usize,
417                    pad_w: pads.get(1).copied().unwrap_or(0).max(0) as usize,
418                },
419            }
420        }
421
422        "AveragePool" => {
423            let kernel = get_attr_ints(node, "kernel_shape").unwrap_or(&[1, 1]);
424            let strides = get_attr_ints(node, "strides").unwrap_or(kernel);
425            match kernel.len() {
426                1 => Op::AvgPool1d {
427                    kernel_l: kernel[0].max(1) as usize,
428                    stride: strides.first().copied().unwrap_or(kernel[0]).max(1) as usize,
429                },
430                3 => Op::AvgPool3d {
431                    kernel_d: kernel[0].max(1) as usize,
432                    kernel_h: kernel[1].max(1) as usize,
433                    kernel_w: kernel[2].max(1) as usize,
434                    stride_d: strides.first().copied().unwrap_or(kernel[0]).max(1) as usize,
435                    stride_h: strides.get(1).copied().unwrap_or(kernel[1]).max(1) as usize,
436                    stride_w: strides.get(2).copied().unwrap_or(kernel[2]).max(1) as usize,
437                },
438                _ => Op::AvgPool2d {
439                    kernel_h: kernel.first().copied().unwrap_or(1).max(1) as usize,
440                    kernel_w: kernel.get(1).copied().unwrap_or(1).max(1) as usize,
441                    stride_h: strides.first().copied().unwrap_or(1).max(1) as usize,
442                    stride_w: strides.get(1).copied().unwrap_or(1).max(1) as usize,
443                },
444            }
445        }
446
447        "GlobalAveragePool" => Op::GlobalAvgPool,
448        "GlobalMaxPool" => Op::GlobalMaxPool,
449
450        "LpPool" => {
451            let kernel = get_attr_ints(node, "kernel_shape").unwrap_or(&[1, 1]);
452            let strides = get_attr_ints(node, "strides").unwrap_or(kernel);
453            let p = get_attr_int(node, "p", 2) as f64;
454            match kernel.len() {
455                1 => Op::LpPool1d {
456                    kernel_l: kernel[0].max(1) as usize,
457                    stride: strides.first().copied().unwrap_or(kernel[0]).max(1) as usize,
458                    p,
459                },
460                3 => Op::LpPool3d {
461                    kernel_d: kernel[0].max(1) as usize,
462                    kernel_h: kernel[1].max(1) as usize,
463                    kernel_w: kernel[2].max(1) as usize,
464                    stride_d: strides.first().copied().unwrap_or(kernel[0]).max(1) as usize,
465                    stride_h: strides.get(1).copied().unwrap_or(kernel[1]).max(1) as usize,
466                    stride_w: strides.get(2).copied().unwrap_or(kernel[2]).max(1) as usize,
467                    p,
468                },
469                _ => Op::LpPool2d {
470                    kernel_h: kernel.first().copied().unwrap_or(1).max(1) as usize,
471                    kernel_w: kernel.get(1).copied().unwrap_or(1).max(1) as usize,
472                    stride_h: strides.first().copied().unwrap_or(1).max(1) as usize,
473                    stride_w: strides.get(1).copied().unwrap_or(1).max(1) as usize,
474                    p,
475                },
476            }
477        }
478
479        "MaxUnpool" => {
480            let kernel = get_attr_ints(node, "kernel_shape").unwrap_or(&[1, 1]);
481            let strides = get_attr_ints(node, "strides").unwrap_or(kernel);
482            Op::MaxUnpool {
483                kernel_h: kernel.first().copied().unwrap_or(1).max(1) as usize,
484                kernel_w: kernel.get(1).copied().unwrap_or(1).max(1) as usize,
485                stride_h: strides.first().copied().unwrap_or(1).max(1) as usize,
486                stride_w: strides.get(1).copied().unwrap_or(1).max(1) as usize,
487            }
488        }
489
490        // ---------------------------------------------------------------
491        // Normalisation
492        // ---------------------------------------------------------------
493        "BatchNormalization" => {
494            let eps = get_attr_float(node, "epsilon", 1e-5) as f64;
495            let momentum = get_attr_float(node, "momentum", 0.9) as f64;
496            // Use BatchNorm2d as the representative variant (most common in practice).
497            Op::BatchNorm2d {
498                num_features: 0,
499                eps,
500                momentum,
501                affine: true,
502                track_running_stats: true,
503            }
504        }
505
506        "LayerNormalization" | "LayerNorm" => {
507            let eps = get_attr_float(node, "epsilon", 1e-5) as f64;
508            Op::LayerNorm {
509                normalized_shape: vec![],
510                eps,
511                affine: true,
512            }
513        }
514
515        "RMSNormalization" => {
516            let eps = get_attr_float(node, "epsilon", 1e-5) as f64;
517            Op::RmsNorm {
518                normalized_shape: vec![],
519                eps,
520                affine: true,
521            }
522        }
523
524        "GroupNormalization" => {
525            let num_groups = get_attr_int(node, "num_groups", 1).max(1) as usize;
526            let eps = get_attr_float(node, "epsilon", 1e-5) as f64;
527            Op::GroupNorm {
528                num_groups,
529                num_channels: 0,
530                eps,
531                affine: true,
532            }
533        }
534
535        "InstanceNormalization" => {
536            let eps = get_attr_float(node, "epsilon", 1e-5) as f64;
537            Op::InstanceNorm2d {
538                num_features: 0,
539                eps,
540                momentum: 0.1,
541                affine: true,
542                track_running_stats: false,
543            }
544        }
545
546        "LRN" => {
547            let alpha = get_attr_float(node, "alpha", 0.0001) as f64;
548            let beta = get_attr_float(node, "beta", 0.75) as f64;
549            let bias = get_attr_float(node, "bias", 1.0) as f64;
550            let size = get_attr_int(node, "size", 1).max(1) as usize;
551            Op::LRN {
552                alpha,
553                beta,
554                bias,
555                size,
556            }
557        }
558
559        "MeanVarianceNormalization" => {
560            let axes = get_attr_ints(node, "axes").unwrap_or(&[0, 2, 3]).to_vec();
561            Op::MeanVarianceNormalization { axes }
562        }
563
564        "LpNormalization" => {
565            let axis = get_attr_int(node, "axis", -1);
566            let p = get_attr_int(node, "p", 2);
567            Op::LpNormalization { axis, p }
568        }
569
570        // ---------------------------------------------------------------
571        // Activations
572        // ---------------------------------------------------------------
573        "Elu" => Op::Elu {
574            alpha: get_attr_float(node, "alpha", 1.0) as f64,
575        },
576        "Selu" => Op::Selu,
577        "Celu" => Op::Celu {
578            alpha: get_attr_float(node, "alpha", 1.0) as f64,
579        },
580        "Gelu" => Op::Gelu,
581        "Mish" => Op::Mish,
582        "HardSigmoid" => Op::Hardsigmoid,
583        "HardSwish" => Op::Hardswish,
584        "Swish" => Op::Swish,
585        "LeakyRelu" => Op::LeakyRelu {
586            negative_slope: get_attr_float(node, "alpha", 0.01) as f64,
587        },
588        "Softplus" => Op::Softplus {
589            beta: 1.0,
590            threshold: 20.0,
591        },
592        "Softsign" => Op::Softsign,
593        "LogSoftmax" => Op::LogSoftmax {
594            axis: get_attr_int(node, "axis", -1),
595        },
596        "Hardmax" => Op::Hardmax {
597            axis: get_attr_int(node, "axis", -1),
598        },
599        "PRelu" => Op::PRelu,
600        "ThresholdedRelu" => Op::ThresholdedRelu {
601            alpha: get_attr_float(node, "alpha", 1.0) as f64,
602        },
603        "Shrink" => Op::Shrink {
604            lambd: get_attr_float(node, "lambd", 0.5) as f64,
605            bias: get_attr_float(node, "bias", 0.0) as f64,
606        },
607        "Clip" => Op::Clip,
608
609        // ---------------------------------------------------------------
610        // Element-wise unary
611        // ---------------------------------------------------------------
612        "Abs" => Op::Abs,
613        "Neg" => Op::Neg,
614        "Ceil" => Op::Ceil,
615        "Floor" => Op::Floor,
616        "Round" => Op::Round,
617        "Sqrt" => Op::Sqrt,
618        "Reciprocal" => Op::Reciprocal,
619        "Exp" => Op::Exp,
620        "Log" => Op::Log,
621        "Erf" => Op::Erf,
622        "Sign" => Op::Sign,
623        "IsNaN" => Op::IsNaN,
624        "IsInf" => Op::IsInf {
625            detect_negative: get_attr_int(node, "detect_negative", 1) != 0,
626            detect_positive: get_attr_int(node, "detect_positive", 1) != 0,
627        },
628        "Not" => Op::Not,
629        "BitwiseNot" => Op::BitwiseNot,
630        "Sin" => Op::Sin,
631        "Cos" => Op::Cos,
632        "Tan" => Op::Tan,
633        "Asin" => Op::Asin,
634        "Acos" => Op::Acos,
635        "Atan" => Op::Atan,
636        "Sinh" => Op::Sinh,
637        "Cosh" => Op::Cosh,
638        "Asinh" => Op::Asinh,
639        "Acosh" => Op::Acosh,
640        "Atanh" => Op::Atanh,
641
642        // ---------------------------------------------------------------
643        // Element-wise binary / variadic
644        // ---------------------------------------------------------------
645        "Mul" => Op::Mul,
646        "Sub" => Op::Sub,
647        "Div" => Op::Div,
648        "Pow" => Op::Pow,
649        "Mod" => Op::Mod {
650            fmod: get_attr_int(node, "fmod", 0) != 0,
651        },
652        "Min" => Op::ElemMin,
653        "Max" => Op::ElemMax,
654        "Mean" => Op::ElemMean,
655        "Sum" => Op::ElemSum,
656        "Equal" => Op::Equal,
657        "Greater" => Op::Greater,
658        "GreaterOrEqual" => Op::GreaterOrEqual,
659        "Less" => Op::Less,
660        "LessOrEqual" => Op::LessOrEqual,
661        "And" => Op::And,
662        "Or" => Op::Or,
663        "Xor" => Op::Xor,
664        "BitwiseAnd" => Op::BitwiseAnd,
665        "BitwiseOr" => Op::BitwiseOr,
666        "BitwiseXor" => Op::BitwiseXor,
667        "BitShift" => Op::BitShift {
668            direction: get_attr_string(node, "direction", "LEFT").to_string(),
669        },
670
671        // ---------------------------------------------------------------
672        // Tensor structural
673        // ---------------------------------------------------------------
674        "Reshape" => Op::Reshape,
675        "Transpose" => Op::Transpose {
676            perm: get_attr_ints(node, "perm")
677                .map(|v| v.iter().map(|&i| i.max(0) as usize).collect())
678                .unwrap_or_default(),
679        },
680        "Squeeze" => Op::Squeeze {
681            axes: get_attr_ints(node, "axes").unwrap_or_default().to_vec(),
682        },
683        "Unsqueeze" => Op::Unsqueeze {
684            axes: get_attr_ints(node, "axes").unwrap_or_default().to_vec(),
685        },
686        "Concat" => Op::Concat {
687            axis: get_attr_int(node, "axis", 0),
688        },
689        "Split" => Op::Split {
690            axis: get_attr_int(node, "axis", 0),
691            num_outputs: node.output.len().max(1),
692        },
693        "Slice" => Op::Slice,
694        "Gather" => Op::Gather {
695            axis: get_attr_int(node, "axis", 0),
696        },
697        "GatherElements" => Op::GatherElements {
698            axis: get_attr_int(node, "axis", 0),
699        },
700        "GatherND" => Op::GatherND {
701            batch_dims: get_attr_int(node, "batch_dims", 0),
702        },
703        "ScatterElements" => Op::ScatterElements {
704            axis: get_attr_int(node, "axis", 0),
705        },
706        "ScatterND" => Op::ScatterND,
707        "Scatter" => Op::Scatter {
708            axis: get_attr_int(node, "axis", 0),
709        },
710        "TensorScatter" => Op::TensorScatter,
711        "Tile" => Op::Tile,
712        "Expand" => Op::Expand,
713        "Shape" => Op::ShapeOf {
714            start: get_attr_int(node, "start", 0),
715            end: get_attr_int(node, "end", i64::MAX),
716        },
717        "Size" => Op::SizeOf,
718        "Identity" => Op::Identity,
719        "Cast" => {
720            let to = get_attr_int(node, "to", 1);
721            let dtype = tensor_dtype_to_dtype(to as i32)?;
722            Op::Cast { to: dtype }
723        }
724        "CastLike" => Op::CastLike,
725        "Where" => Op::Where,
726        "Compress" => Op::Compress {
727            axis: get_attr_int(node, "axis", 0),
728        },
729        "Range" => Op::Range,
730        "Constant" => {
731            // Shape/dtype extracted from the tensor attribute.
732            if let Some(attr) = node.attribute.iter().find(|a| a.name == "value") {
733                let tensor = &attr.t;
734                let dtype = tensor_dtype_to_dtype(tensor.data_type)?;
735                let shape = dims_to_shape(&tensor.dims);
736                Op::Constant { dtype, shape }
737            } else {
738                Op::Constant {
739                    dtype: DtypeRepr::F32,
740                    shape: vec![],
741                }
742            }
743        }
744        "ConstantOfShape" => {
745            let dtype = if let Some(attr) = node.attribute.iter().find(|a| a.name == "value") {
746                tensor_dtype_to_dtype(attr.t.data_type)?
747            } else {
748                DtypeRepr::F32
749            };
750            Op::ConstantOfShape { dtype }
751        }
752        "Trilu" => Op::Trilu {
753            upper: get_attr_int(node, "upper", 1) != 0,
754        },
755        "BitCast" => {
756            let to = get_attr_int(node, "to", 1);
757            let dtype = tensor_dtype_to_dtype(to as i32)?;
758            Op::BitCast { to: dtype }
759        }
760        "Pad" => Op::Pad {
761            mode: get_attr_string(node, "mode", "constant").to_string(),
762        },
763        "ReverseSequence" => Op::ReverseSequence {
764            batch_axis: get_attr_int(node, "batch_axis", 1),
765            time_axis: get_attr_int(node, "time_axis", 0),
766        },
767        "NonZero" => Op::NonZero,
768
769        // ---------------------------------------------------------------
770        // Matrix
771        // ---------------------------------------------------------------
772        "Gemm" => Op::Gemm {
773            alpha: get_attr_float(node, "alpha", 1.0) as f64,
774            beta: get_attr_float(node, "beta", 1.0) as f64,
775            trans_a: get_attr_int(node, "transA", 0) != 0,
776            trans_b: get_attr_int(node, "transB", 0) != 0,
777        },
778        "MatMul" => Op::MatMul,
779        "MatMulInteger" => Op::MatMulInteger,
780        "Einsum" => Op::Einsum {
781            equation: get_attr_string(node, "equation", "").to_string(),
782        },
783        "Det" => Op::Det,
784        "QLinearMatMul" => Op::QLinearMatMul,
785        "ConvInteger" => Op::ConvInteger {
786            groups: get_attr_int(node, "group", 1).max(1) as usize,
787        },
788        "DeformConv" => Op::DeformConv {
789            group: get_attr_int(node, "group", 1).max(1) as usize,
790            offset_group: get_attr_int(node, "offset_group", 1).max(1) as usize,
791        },
792        "QLinearConv" => Op::QLinearConv {
793            groups: get_attr_int(node, "group", 1).max(1) as usize,
794        },
795        "Col2Im" => {
796            let kernel = get_attr_ints(node, "kernel_shape").unwrap_or(&[1, 1]);
797            Op::Col2Im {
798                kernel_h: kernel.first().copied().unwrap_or(1).max(1) as usize,
799                kernel_w: kernel.get(1).copied().unwrap_or(1).max(1) as usize,
800            }
801        }
802
803        // ---------------------------------------------------------------
804        // Reductions
805        // ---------------------------------------------------------------
806        "ReduceSum" => Op::ReduceSum {
807            keepdims: get_attr_int(node, "keepdims", 1) != 0,
808            noop_with_empty_axes: get_attr_int(node, "noop_with_empty_axes", 0) != 0,
809        },
810        "ReduceMean" => Op::ReduceMean {
811            keepdims: get_attr_int(node, "keepdims", 1) != 0,
812            noop_with_empty_axes: get_attr_int(node, "noop_with_empty_axes", 0) != 0,
813        },
814        "ReduceMax" => Op::ReduceMax {
815            keepdims: get_attr_int(node, "keepdims", 1) != 0,
816            noop_with_empty_axes: get_attr_int(node, "noop_with_empty_axes", 0) != 0,
817        },
818        "ReduceMin" => Op::ReduceMin {
819            keepdims: get_attr_int(node, "keepdims", 1) != 0,
820            noop_with_empty_axes: get_attr_int(node, "noop_with_empty_axes", 0) != 0,
821        },
822        "ReduceProd" => Op::ReduceProd {
823            keepdims: get_attr_int(node, "keepdims", 1) != 0,
824            noop_with_empty_axes: get_attr_int(node, "noop_with_empty_axes", 0) != 0,
825        },
826        "ReduceL1" => Op::ReduceL1 {
827            keepdims: get_attr_int(node, "keepdims", 1) != 0,
828            noop_with_empty_axes: get_attr_int(node, "noop_with_empty_axes", 0) != 0,
829        },
830        "ReduceL2" => Op::ReduceL2 {
831            keepdims: get_attr_int(node, "keepdims", 1) != 0,
832            noop_with_empty_axes: get_attr_int(node, "noop_with_empty_axes", 0) != 0,
833        },
834        "ReduceLogSum" => Op::ReduceLogSum {
835            keepdims: get_attr_int(node, "keepdims", 1) != 0,
836            noop_with_empty_axes: get_attr_int(node, "noop_with_empty_axes", 0) != 0,
837        },
838        "ReduceLogSumExp" => Op::ReduceLogSumExp {
839            keepdims: get_attr_int(node, "keepdims", 1) != 0,
840            noop_with_empty_axes: get_attr_int(node, "noop_with_empty_axes", 0) != 0,
841        },
842        "ReduceSumSquare" => Op::ReduceSumSquare {
843            keepdims: get_attr_int(node, "keepdims", 1) != 0,
844            noop_with_empty_axes: get_attr_int(node, "noop_with_empty_axes", 0) != 0,
845        },
846        "CumSum" => Op::CumSum {
847            exclusive: get_attr_int(node, "exclusive", 0) != 0,
848            reverse: get_attr_int(node, "reverse", 0) != 0,
849        },
850        "CumProd" => Op::CumProd {
851            exclusive: get_attr_int(node, "exclusive", 0) != 0,
852            reverse: get_attr_int(node, "reverse", 0) != 0,
853        },
854        "ArgMax" => Op::ArgMax {
855            axis: get_attr_int(node, "axis", 0),
856            keepdims: get_attr_int(node, "keepdims", 1) != 0,
857            select_last_index: get_attr_int(node, "select_last_index", 0) != 0,
858        },
859        "ArgMin" => Op::ArgMin {
860            axis: get_attr_int(node, "axis", 0),
861            keepdims: get_attr_int(node, "keepdims", 1) != 0,
862            select_last_index: get_attr_int(node, "select_last_index", 0) != 0,
863        },
864
865        // ---------------------------------------------------------------
866        // Resize / spatial
867        // ---------------------------------------------------------------
868        "Upsample" | "Resize" => Op::Resize {
869            mode: get_attr_string(node, "mode", "nearest").to_string(),
870            coordinate_transformation_mode: get_attr_string(
871                node,
872                "coordinate_transformation_mode",
873                "asymmetric",
874            )
875            .to_string(),
876            antialias: get_attr_int(node, "antialias", 0) != 0,
877        },
878
879        "GridSample" => Op::GridSample {
880            mode: get_attr_string(node, "mode", "bilinear").to_string(),
881            padding_mode: get_attr_string(node, "padding_mode", "zeros").to_string(),
882            align_corners: get_attr_int(node, "align_corners", 0) != 0,
883        },
884
885        "SpaceToDepth" => Op::SpaceToDepth {
886            blocksize: get_attr_int(node, "blocksize", 2).max(1) as usize,
887        },
888        "DepthToSpace" => Op::DepthToSpace {
889            blocksize: get_attr_int(node, "blocksize", 2).max(1) as usize,
890            mode: get_attr_string(node, "mode", "DCR").to_string(),
891        },
892
893        "RoiAlign" => Op::RoiAlign {
894            output_h: get_attr_int(node, "output_height", 1).max(1) as usize,
895            output_w: get_attr_int(node, "output_width", 1).max(1) as usize,
896            sampling_ratio: get_attr_int(node, "sampling_ratio", 0),
897            spatial_scale: get_attr_float(node, "spatial_scale", 1.0) as f64,
898        },
899
900        "AffineGrid" => Op::AffineGrid {
901            align_corners: get_attr_int(node, "align_corners", 0) != 0,
902        },
903
904        "CenterCropPad" => Op::CenterCropPad {
905            axes: get_attr_ints(node, "axes").unwrap_or_default().to_vec(),
906        },
907
908        "NonMaxSuppression" => Op::NonMaxSuppression {
909            center_point_box: get_attr_int(node, "center_point_box", 0) != 0,
910        },
911
912        // ---------------------------------------------------------------
913        // Recurrent
914        // ---------------------------------------------------------------
915        "LSTM" => {
916            let hidden_size = get_attr_int(node, "hidden_size", 1).max(1) as usize;
917            let direction = get_attr_string(node, "direction", "forward").to_string();
918            let bidirectional = direction == "bidirectional";
919            Op::Lstm {
920                hidden_size,
921                direction,
922                bidirectional,
923            }
924        }
925        "GRU" => {
926            let hidden_size = get_attr_int(node, "hidden_size", 1).max(1) as usize;
927            let direction = get_attr_string(node, "direction", "forward").to_string();
928            let bidirectional = direction == "bidirectional";
929            Op::Gru {
930                hidden_size,
931                direction,
932                bidirectional,
933            }
934        }
935        "RNN" => {
936            let hidden_size = get_attr_int(node, "hidden_size", 1).max(1) as usize;
937            let direction = get_attr_string(node, "direction", "forward").to_string();
938            let bidirectional = direction == "bidirectional";
939            Op::Rnn {
940                hidden_size,
941                direction,
942                bidirectional,
943            }
944        }
945
946        // ---------------------------------------------------------------
947        // Attention
948        // ---------------------------------------------------------------
949        "Attention" => Op::MultiHeadAttention {
950            q_num_heads: get_attr_int(node, "q_num_heads", 1).max(1) as usize,
951            kv_num_heads: get_attr_int(node, "kv_num_heads", 1).max(1) as usize,
952        },
953        "RotaryEmbedding" => Op::RotaryEmbedding,
954        // score_mod/prob_mod are GRAPH-typed attributes (a user-provided scoring subgraph) --
955        // not captured, consistent with Loop/If/Scan not capturing their subgraph bodies.
956        "FlexAttention" => Op::FlexAttention {
957            scale: get_attr_float(node, "scale", 0.0) as f64,
958        },
959        "LinearAttention" => Op::LinearAttention {
960            q_num_heads: get_attr_int(node, "q_num_heads", 1).max(1) as usize,
961            kv_num_heads: get_attr_int(node, "kv_num_heads", 1).max(1) as usize,
962            update_rule: get_attr_string(node, "update_rule", "").to_string(),
963            scale: get_attr_float(node, "scale", 0.0) as f64,
964        },
965        "CausalConvWithState" => Op::CausalConvWithState {
966            activation: get_attr_string(node, "activation", "").to_string(),
967        },
968
969        // ---------------------------------------------------------------
970        // Misc
971        // ---------------------------------------------------------------
972        "TopK" => Op::TopK {
973            axis: get_attr_int(node, "axis", -1),
974            largest: get_attr_int(node, "largest", 1) != 0,
975            sorted: get_attr_int(node, "sorted", 1) != 0,
976        },
977        "Unique" => Op::Unique {
978            sorted: get_attr_int(node, "sorted", 1) != 0,
979        },
980        "Dropout" => Op::Dropout {
981            training_mode: get_attr_int(node, "training_mode", 0) != 0,
982        },
983        "EyeLike" => {
984            let to = get_attr_int(node, "dtype", -1);
985            let dtype = if to >= 0 {
986                Some(tensor_dtype_to_dtype(to as i32)?)
987            } else {
988                None
989            };
990            Op::EyeLike {
991                dtype,
992                k: get_attr_int(node, "k", 0),
993            }
994        }
995        "OneHot" => Op::OneHot {
996            axis: get_attr_int(node, "axis", -1),
997        },
998        "Bernoulli" => {
999            let to = get_attr_int(node, "dtype", -1);
1000            let dtype = if to >= 0 {
1001                Some(tensor_dtype_to_dtype(to as i32)?)
1002            } else {
1003                None
1004            };
1005            Op::Bernoulli { dtype }
1006        }
1007        "RandomUniformLike" => {
1008            let to = get_attr_int(node, "dtype", -1);
1009            let dtype = if to >= 0 {
1010                Some(tensor_dtype_to_dtype(to as i32)?)
1011            } else {
1012                None
1013            };
1014            Op::RandomUniformLike {
1015                dtype,
1016                high: get_attr_float(node, "high", 1.0) as f64,
1017                low: get_attr_float(node, "low", 0.0) as f64,
1018            }
1019        }
1020
1021        // ---------------------------------------------------------------
1022        // Quantisation
1023        // ---------------------------------------------------------------
1024        "QuantizeLinear" => Op::QuantizeLinear {
1025            axis: get_attr_int(node, "axis", 1),
1026            saturate: get_attr_int(node, "saturate", 1) != 0,
1027        },
1028        "DequantizeLinear" => Op::DequantizeLinear {
1029            axis: get_attr_int(node, "axis", 1),
1030        },
1031        "DynamicQuantizeLinear" => Op::DynamicQuantizeLinear,
1032
1033        // ---------------------------------------------------------------
1034        // Signal
1035        // ---------------------------------------------------------------
1036        "DFT" => Op::Dft {
1037            inverse: get_attr_int(node, "inverse", 0) != 0,
1038            onesided: get_attr_int(node, "onesided", 0) != 0,
1039        },
1040        "STFT" => Op::Stft,
1041        "MelWeightMatrix" => Op::MelWeightMatrix,
1042        "HannWindow" => Op::HannWindow {
1043            periodic: get_attr_int(node, "periodic", 1) != 0,
1044        },
1045        "BlackmanWindow" => Op::BlackmanWindow {
1046            periodic: get_attr_int(node, "periodic", 1) != 0,
1047        },
1048        "HammingWindow" => Op::HammingWindow {
1049            periodic: get_attr_int(node, "periodic", 1) != 0,
1050        },
1051
1052        // ---------------------------------------------------------------
1053        // Loss
1054        // ---------------------------------------------------------------
1055        "NegativeLogLikelihoodLoss" => Op::NegativeLogLikelihoodLoss {
1056            reduction: get_attr_string(node, "reduction", "mean").to_string(),
1057        },
1058        "SoftmaxCrossEntropyLoss" => Op::SoftmaxCrossEntropyLoss {
1059            reduction: get_attr_string(node, "reduction", "mean").to_string(),
1060        },
1061
1062        // ---------------------------------------------------------------
1063        // Sequences / optionals
1064        // ---------------------------------------------------------------
1065        "SequenceAt" => Op::SequenceAt,
1066        "SequenceConstruct" => Op::SequenceConstruct,
1067        "SequenceEmpty" => Op::SequenceEmpty,
1068        "SequenceErase" => Op::SequenceErase,
1069        "SequenceInsert" => Op::SequenceInsert,
1070        "SequenceLength" => Op::SequenceLength,
1071        "SequenceMap" => Op::SequenceMap,
1072        "SplitToSequence" => Op::SplitToSequence {
1073            axis: get_attr_int(node, "axis", 0),
1074            keepdims: get_attr_int(node, "keepdims", 1) != 0,
1075        },
1076        "ConcatFromSequence" => Op::ConcatFromSequence {
1077            axis: get_attr_int(node, "axis", 0),
1078            new_axis: get_attr_int(node, "new_axis", 0) != 0,
1079        },
1080        "OptionalGetElement" => Op::OptionalGetElement,
1081        "OptionalHasElement" => Op::OptionalHasElement,
1082
1083        // ---------------------------------------------------------------
1084        // Control flow
1085        // ---------------------------------------------------------------
1086        "Loop" => Op::Loop,
1087        "Scan" => Op::Scan {
1088            num_scan_inputs: get_attr_int(node, "num_scan_inputs", 0),
1089        },
1090        "If" => Op::If,
1091
1092        // ---------------------------------------------------------------
1093        // Optimiser ops
1094        // ---------------------------------------------------------------
1095        "Adagrad" => Op::Adagrad,
1096        "Adam" => Op::Adam,
1097        "Momentum" => Op::Momentum,
1098        "Gradient" => Op::Gradient,
1099
1100        // ---------------------------------------------------------------
1101        // String / NLP
1102        // ---------------------------------------------------------------
1103        "StringNormalizer" => Op::StringNormalizer,
1104        "RegexFullMatch" => Op::RegexFullMatch {
1105            pattern: get_attr_string(node, "pattern", "").to_string(),
1106        },
1107        "StringConcat" => Op::StringConcat,
1108        "StringSplit" => Op::StringSplit,
1109        "TfIdfVectorizer" => Op::TfIdfVectorizer,
1110        "LabelEncoder" => Op::LabelEncoder,
1111
1112        // ---------------------------------------------------------------
1113        // Other ML / deprecated
1114        // ---------------------------------------------------------------
1115        "ArrayFeatureExtractor" => Op::ArrayFeatureExtractor,
1116        "Binarizer" => Op::Binarizer {
1117            threshold: get_attr_float(node, "threshold", 0.0) as f64,
1118        },
1119        "TreeEnsemble" => Op::TreeEnsemble,
1120        "ImageDecoder" => Op::ImageDecoder,
1121
1122        other => {
1123            return Err(Error::InvalidModel(format!(
1124                "Unsupported ONNX op '{}' in node '{}'",
1125                other, node.name
1126            ))
1127            .into());
1128        }
1129    };
1130    Ok(op)
1131}
teeny_onnx/onnx.rs

teeny_onnx/
onnx.rs
