Fix nondeterministic Loop output binding with multiple state variables

onnxscript/_internal/converter.py

@@ -1234,9 +1234,12 @@ def _translate_loop_stmt(self, loop_stmt: ast.For | ast.While) -> None:
         vars_def_in_loop = self.analyzer.assigned_vars(loop_stmt.body)
         live_out = self.analyzer.live_out(loop_stmt)
         assert live_out is not None, "live_out cannot be None here."
-        loop_state_vars = vars_def_in_loop.intersection(exposed_uses | live_out)
-        scan_outputs = set()  # TODO
-        outputs = list(loop_state_vars | scan_outputs)
+        # A list, not a set: this order is used to build both the Loop node's
+        # loop-carried inputs (below) and its output names (via `outputs`), and ONNX
+        # matches Loop inputs to outputs positionally, not by name.
+        loop_state_vars = list(vars_def_in_loop.intersection(exposed_uses | live_out))
+        scan_outputs = []  # TODO
+        outputs = loop_state_vars + scan_outputs
 
         # loop-condition:
         # o_loop_condition = self._emit_const(True, "true", self._source_of(loop_stmt))

tests/loop_test.py

@@ -1,5 +1,9 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
+import os
+import subprocess
+import sys
+import textwrap
 import unittest
 
 import numpy as np
@@ -9,6 +13,33 @@
 from onnxscript.onnx_types import FLOAT, INT64
 from tests.common import testutils
 
+# tests/models/loop_multi_state.py has 5 loop-carried variables: enough for the
+# converter's output-name ordering (built via a set union) to diverge from its
+# input ordering (built by iterating the same set directly) under CPython's
+# hash randomization. Run as a subprocess under a fixed PYTHONHASHSEED so the
+# divergence reproduces deterministically instead of depending on whatever
+# seed the test runner's own process happens to pick.
+_MULTI_STATE_LOOP_CHECK = textwrap.dedent(
+    """
+    import numpy as np
+    import onnx
+    from onnxruntime import InferenceSession
+
+    from tests.models.loop_multi_state import loop_multi_state
+
+    x = np.array([2.0, 3.0], dtype=np.float32)
+    eager = loop_multi_state(x, 3)
+
+    model = loop_multi_state.to_model_proto()
+    onnx.checker.check_model(model)
+    sess = InferenceSession(model.SerializeToString(), providers=("CPUExecutionProvider",))
+    got = sess.run(None, {"x": x, "n": np.array(3, dtype=np.int64)})
+
+    for eager_arr, got_arr in zip(eager, got):
+        assert np.allclose(eager_arr, got_arr), (eager_arr, got_arr)
+    """
+)
+
 
 class LoopOpTest(testutils.TestBase):
     def test_loop(self):
@@ -44,6 +75,23 @@ def sumprod(x: FLOAT["N"], N: INT64) -> (FLOAT["N"], FLOAT["N"]):  # noqa: F821
 
         self.validate(sumprod)
 
+    def test_loop_state_var_output_order_matches_eager(self):
+        """A Loop with >1 loop-carried variable must bind each returned Python
+        name to the ONNX value it actually computes, not to whichever value a
+        hash-randomized set iteration happened to line up with it.
+        """
+        for seed in ("3", "4", "8"):
+            env = dict(os.environ, PYTHONHASHSEED=seed)
+            result = subprocess.run(
+                [sys.executable, "-c", _MULTI_STATE_LOOP_CHECK],
+                capture_output=True,
+                text=True,
+                env=env,
+                cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
+            )
+            with self.subTest(seed=seed):
+                self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
+
 
 if __name__ == "__main__":
     unittest.main()

tests/models/loop_multi_state.py

@@ -0,0 +1,33 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from onnxscript import script
+from onnxscript.onnx_opset import opset15 as op
+from onnxscript.onnx_types import FLOAT, INT64
+
+# Five loop-carried state variables: enough for the converter's output-name
+# ordering (built via a set union) to diverge from its input ordering (built
+# by iterating the same set directly) under CPython's hash randomization.
+
+
+@script()
+def loop_multi_state(
+    x: FLOAT["N"],
+    n: INT64,  # noqa: F821
+) -> (FLOAT["N"], FLOAT["N"], FLOAT["N"], FLOAT["N"], FLOAT["N"]):  # noqa: F821
+    a = op.Identity(x)
+    b = op.Identity(x)
+    c = op.Identity(x)
+    d = op.Identity(x)
+    e = op.Identity(x)
+    i = 0
+    cond = True
+    while cond:
+        a = op.Add(a, x)
+        b = op.Mul(b, x)
+        c = op.Sub(c, x)
+        d = op.Div(d, x)
+        e = op.Add(e, b)
+        i = i + 1
+        cond = i < n
+    return a, b, c, d, e