Skip to main content

onnx_extractor/
operation.rs

1use crate::AttributeValue;
2use std::collections::HashMap;
3
4/// An ONNX operation/node in the computational graph
5#[derive(Debug)]
6pub struct Operation {
7    name: Option<String>,
8    op_type: String,
9    inputs: Vec<String>,
10    outputs: Vec<String>,
11    attributes: HashMap<String, AttributeValue>,
12}
13
14impl Operation {
15    pub(crate) fn new(
16        name: Option<String>,
17        op_type: String,
18        inputs: Vec<String>,
19        outputs: Vec<String>,
20        attributes: HashMap<String, AttributeValue>,
21    ) -> Self {
22        Operation {
23            name,
24            op_type,
25            inputs,
26            outputs,
27            attributes,
28        }
29    }
30
31    /// Operation name
32    pub fn name(&self) -> Option<&str> {
33        self.name.as_deref()
34    }
35
36    /// Operation type (e.g., "Conv", "Relu")
37    pub fn op_type(&self) -> &str {
38        &self.op_type
39    }
40
41    /// Input tensor names
42    pub fn inputs(&self) -> &[String] {
43        &self.inputs
44    }
45
46    /// Output tensor names
47    pub fn outputs(&self) -> &[String] {
48        &self.outputs
49    }
50
51    /// Reference to all attributes
52    pub fn attributes(&self) -> &HashMap<String, AttributeValue> {
53        &self.attributes
54    }
55
56    /// Mutable reference to all attributes.
57    ///
58    /// Allows removing, draining, or modifying attributes directly.
59    pub fn attributes_mut(&mut self) -> &mut HashMap<String, AttributeValue> {
60        &mut self.attributes
61    }
62}