torsh-fx 0.1.2

Graph-based model representation and transformation for ToRSh
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
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

//! Essential FxGraph implementation and core functionality

use crate::codegen::CodeGenerator;
use crate::fx::types::{Edge, Node};
use crate::onnx_export::{export_to_onnx, OnnxExporter, OnnxModel};
use crate::{FxGraph, TorshResult};
use petgraph::graph::{Graph, NodeIndex};

/// FX Graph representation
#[derive(Debug, Clone)]
pub struct FxGraphCore {
    pub graph: Graph<Node, Edge>,
    pub inputs: Vec<NodeIndex>,
    pub outputs: Vec<NodeIndex>,
}

impl FxGraph {
    /// Create new graph
    pub fn new() -> Self {
        Self {
            graph: Graph::new(),
            inputs: Vec::new(),
            outputs: Vec::new(),
        }
    }

    /// Get the number of nodes in the graph
    pub fn node_count(&self) -> usize {
        self.graph.node_count()
    }

    /// Get the number of edges in the graph
    pub fn edge_count(&self) -> usize {
        self.graph.edge_count()
    }

    /// Get input nodes
    pub fn inputs(&self) -> &[NodeIndex] {
        &self.inputs
    }

    /// Get output nodes
    pub fn outputs(&self) -> &[NodeIndex] {
        &self.outputs
    }

    /// Get a node by index
    pub fn get_node(&self, idx: NodeIndex) -> Option<&Node> {
        self.graph.node_weight(idx)
    }

    /// Add a node to the graph
    pub fn add_node(&mut self, node: Node) -> NodeIndex {
        self.graph.add_node(node)
    }

    /// Add an edge to the graph
    pub fn add_edge(&mut self, source: NodeIndex, target: NodeIndex, edge: Edge) {
        self.graph.add_edge(source, target, edge);
    }

    /// Add an input node
    pub fn add_input(&mut self, input: NodeIndex) {
        self.inputs.push(input);
    }

    /// Add an output node
    pub fn add_output(&mut self, output: NodeIndex) {
        self.outputs.push(output);
    }

    /// Iterate over all nodes
    pub fn nodes(&self) -> impl Iterator<Item = (NodeIndex, &Node)> {
        self.graph
            .node_indices()
            .map(move |idx| (idx, &self.graph[idx]))
    }

    /// Print the graph structure
    pub fn print(&self) {
        println!("FX Graph:");
        println!("  Nodes: {}", self.node_count());
        println!("  Edges: {}", self.edge_count());
        println!("  Inputs: {:?}", self.inputs);
        println!("  Outputs: {:?}", self.outputs);

        for (idx, node) in self.nodes() {
            println!("  Node {:?}: {:?}", idx, node);
        }
    }

    /// Generate code for the graph using the specified target language
    pub fn generate_code(&self, target: &str) -> TorshResult<String> {
        let generator = CodeGenerator::new();
        generator.generate_code(self, target)
    }

    /// Generate Python code for the graph
    pub fn to_python(&self) -> TorshResult<String> {
        self.generate_code("python")
    }

    /// Generate C++ code for the graph
    pub fn to_cpp(&self) -> TorshResult<String> {
        self.generate_code("cpp")
    }

    /// Export the graph to ONNX format
    pub fn to_onnx(&self) -> TorshResult<OnnxModel> {
        export_to_onnx(self, None)
    }

    /// Export the graph to ONNX format with a custom model name
    pub fn to_onnx_named(&self, model_name: String) -> TorshResult<OnnxModel> {
        export_to_onnx(self, Some(model_name))
    }

    /// Export the graph to ONNX JSON format
    pub fn to_onnx_json(&self) -> TorshResult<String> {
        let exporter = OnnxExporter::new();
        exporter.export_to_json(self)
    }
}

impl Default for FxGraph {
    fn default() -> Self {
        Self::new()
    }
}