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
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155

//! Graph construction utilities and convenience methods

use crate::fx::types::{Edge, Node};
use crate::FxGraph;

impl FxGraph {
    /// Convenience method to create a graph with a single operation
    pub fn single_op(op_name: &str, input_names: Vec<String>) -> Self {
        let mut graph = Self::new();

        // Add input nodes
        let mut input_indices = Vec::new();
        for input_name in &input_names {
            let input_idx = graph.add_node(Node::Input(input_name.clone()));
            input_indices.push(input_idx);
            graph.add_input(input_idx);
        }

        // Add operation node
        let op_idx = graph.add_node(Node::Call(op_name.to_string(), input_names.clone()));

        // Add output node
        let output_idx = graph.add_node(Node::Output);
        graph.add_output(output_idx);

        // Connect inputs to operation
        for (i, input_idx) in input_indices.iter().enumerate() {
            graph.add_edge(
                *input_idx,
                op_idx,
                Edge {
                    name: format!("input_{i}"),
                },
            );
        }

        // Connect operation to output
        graph.add_edge(
            op_idx,
            output_idx,
            Edge {
                name: "output".to_string(),
            },
        );

        graph
    }

    /// Convenience method to chain multiple operations sequentially
    pub fn sequential_ops(ops: &[&str]) -> Self {
        let mut graph = Self::new();

        if ops.is_empty() {
            return graph;
        }

        // Add input node
        let input_idx = graph.add_node(Node::Input("input".to_string()));
        graph.add_input(input_idx);

        let mut prev_idx = input_idx;

        // Add operations sequentially
        for (i, op_name) in ops.iter().enumerate() {
            let op_idx = graph.add_node(Node::Call(
                op_name.to_string(),
                vec![format!("intermediate_{i}")],
            ));

            graph.add_edge(
                prev_idx,
                op_idx,
                Edge {
                    name: format!("connection_{i}"),
                },
            );

            prev_idx = op_idx;
        }

        // Add output node
        let output_idx = graph.add_node(Node::Output);
        graph.add_output(output_idx);

        graph.add_edge(
            prev_idx,
            output_idx,
            Edge {
                name: "final_output".to_string(),
            },
        );

        graph
    }

    /// Create a minimal test graph for debugging
    pub fn debug_minimal() -> Self {
        Self::sequential_ops(&["debug_op"])
    }

    /// Create a test graph with branching for debugging
    pub fn debug_branching() -> Self {
        let mut graph = Self::new();

        // Create branching structure: input -> op1 -> op2
        //                                   -> op3 -> output
        let input = graph.add_node(Node::Input("x".to_string()));
        let op1 = graph.add_node(Node::Call("branch_op1".to_string(), vec!["x".to_string()]));
        let op2 = graph.add_node(Node::Call("branch_op2".to_string(), vec![]));
        let op3 = graph.add_node(Node::Call("branch_op3".to_string(), vec![]));
        let output = graph.add_node(Node::Output);

        graph.add_input(input);
        graph.add_output(output);

        // Connect the branching structure
        graph.add_edge(
            input,
            op1,
            Edge {
                name: "input_to_op1".to_string(),
            },
        );
        graph.add_edge(
            op1,
            op2,
            Edge {
                name: "op1_to_op2".to_string(),
            },
        );
        graph.add_edge(
            op1,
            op3,
            Edge {
                name: "op1_to_op3".to_string(),
            },
        );
        graph.add_edge(
            op2,
            output,
            Edge {
                name: "op2_to_output".to_string(),
            },
        );
        graph.add_edge(
            op3,
            output,
            Edge {
                name: "op3_to_output".to_string(),
            },
        );

        graph
    }
}