Fixes

jafioti · Jan 7, 2024 · 1a1ba52 · 1a1ba52
1 parent 4e5300c
commit 1a1ba52
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 3 deletions.
diff --git a/src/compilers/generic.rs b/src/compilers/generic.rs
@@ -419,7 +419,14 @@ impl Compiler for ArithmeticElimination {
             {
                 continue;
             }
-            move_outgoing_edge(add, x, &mut graph.graph);
+            for output in graph
+                .get_dests(add)
+                .into_iter()
+                .map(|(o, _)| o)
+                .collect::<Vec<_>>()
+            {
+                move_incoming_edge(add, output, &mut graph.graph);
+            }
             move_references(
                 &mut remap,
                 &mut graph.no_delete,
@@ -470,7 +477,14 @@ impl Compiler for ArithmeticElimination {
             {
                 continue;
             }
-            move_outgoing_edge(mul, a, &mut graph.graph);
+            for output in graph
+                .get_dests(mul)
+                .into_iter()
+                .map(|(o, _)| o)
+                .collect::<Vec<_>>()
+            {
+                move_incoming_edge(mul, output, &mut graph.graph);
+            }
             move_references(
                 &mut remap,
                 &mut graph.no_delete,

diff --git a/src/nn/embedding.rs b/src/nn/embedding.rs
@@ -70,7 +70,6 @@ mod tests {
         let mut b = model.forward(a).retrieve();
         let mut batch_out = model.forward(batch).retrieve();
 
-        // Uncommenting these lines causes the test to crash
         cx.compile(GenericCompiler::<()>::default(), (&mut b, &mut batch_out));
 
         cx.execute();