fix tied input-output equivalence checking

diff --git a/src/bigint_presentation_code/compiler_ir2.py b/src/bigint_presentation_code/compiler_ir2.py
index 05dc6f49d93e9b0e5896807251352bd960185e0e..7bdd40b0c8931a82d9ed55a5be8a70b0e4fa4171 100644 (file)
--- a/src/bigint_presentation_code/compiler_ir2.py
+++ b/src/bigint_presentation_code/compiler_ir2.py
@@ -1,4 +1,3 @@
-from collections import defaultdict
 import enum
 from abc import ABCMeta, abstractmethod
 from enum import Enum, unique
@@ -941,7 +940,9 @@ class GenericOpProperties:
                 if out.tied_input_index >= len(self.inputs):
                     raise ValueError(f"tied_input_index out of range: {out}")
                 tied_inp = self.inputs[out.tied_input_index]
-                if tied_inp.tied_to_input(out.tied_input_index) != out:
+                expected_out = tied_inp.tied_to_input(out.tied_input_index) \
+                    .with_write_stage(out.write_stage)
+                if expected_out != out:
                     raise ValueError(f"output can't be tied to non-equivalent "
                                      f"input: {out} tied to {tied_inp}")
             if out.fixed_loc is not None: