From 9a194ad6900165f0aa1fe866231aa99f4b889235 Mon Sep 17 00:00:00 2001
From: Felix Bauer <felix.bauer@synsense.ai>
Date: Wed, 13 Nov 2024 09:49:01 +0100
Subject: [PATCH] Fix non-deterministic dynapcnn-network test

---
 tests/test_dynapcnnnetwork/model_dummy_4.py   | 40 +++++++++++++++++--
 .../test_dynapcnnnetwork.py                   | 10 +++++
 2 files changed, 46 insertions(+), 4 deletions(-)

diff --git a/tests/test_dynapcnnnetwork/model_dummy_4.py b/tests/test_dynapcnnnetwork/model_dummy_4.py
index c467dba8..b47120ce 100644
--- a/tests/test_dynapcnnnetwork/model_dummy_4.py
+++ b/tests/test_dynapcnnnetwork/model_dummy_4.py
@@ -117,10 +117,6 @@ def forward(self, x):
 
 snn = SNN(batch_size)
 
-# TODO: This test sometimes fails because the layer that has ID 1
-# sometimes gets ID 2, and the layer with ID 3 gets ID 4
-# This is not a bug in sinabs itself but an issue with the test, becuase
-# the IDs that the layers are assigned do not always have to be the same.
 expected_output = {
     "dcnnl_edges": {
         (0, 1),
@@ -152,3 +148,39 @@ def forward(self, x):
     "output_shape": torch.Size([2, 10, 1, 1]),
     "entry_points": {0},
 }
+
+# Sometimes the layer that usually gets assgined ID1, gets ID2, and the
+# layer with ID 3 gets ID 4. Therefore an alternative solution is defined.
+# This is not a bug in sinabs itself but an issue with the test, becuase
+# the IDs that the layers are assigned do not always have to be the same.
+expected_output["alternative"] = {
+    "dcnnl_edges": {
+        (0, 1),
+        (0, 2),
+        (1, 3),
+        (1, 4),
+        (2, 4),
+        (3, 5),
+        (4, 5),
+        ("input", 0),
+    },
+    "node_source_map": {
+        0: {"input"},
+        1: {0},
+        2: {0},
+        3: {1},
+        4: {1, 2},
+        5: {3, 4},
+    },
+    "destination_map": {
+        0: {1, 2},
+        1: {3, 4},
+        2: {4},
+        3: {5},
+        4: {5},
+        5: {-1},
+    },
+    "sorted_nodes": [0, 1, 2, 3, 4, 5],
+    "output_shape": torch.Size([2, 10, 1, 1]),
+    "entry_points": {0},
+}
diff --git a/tests/test_dynapcnnnetwork/test_dynapcnnnetwork.py b/tests/test_dynapcnnnetwork/test_dynapcnnnetwork.py
index 4c50bc39..293ee451 100644
--- a/tests/test_dynapcnnnetwork/test_dynapcnnnetwork.py
+++ b/tests/test_dynapcnnnetwork/test_dynapcnnnetwork.py
@@ -26,6 +26,16 @@ def test_DynapcnnNetwork(snn, input_shape, batch_size, expected_output):
     output = dcnnnet(x)
 
     module = dcnnnet.dynapcnn_module
+    # For some models there are multiple possible topological sortings,
+    # such that the assigned node IDs are not always the same.
+    # To prevent the following tests from failing, alternative expected
+    # outputs are defined which correspond to different assigned IDs.
+    if (
+        expected_output["dcnnl_edges"] != module._dynapcnnlayer_edges
+        and "alternative" in expected_output
+    ):
+        expected_output = expected_output["alternative"]
+        print("Using algernative node ID assignment")
     assert (
         expected_output["dcnnl_edges"] == module._dynapcnnlayer_edges
     ), "wrong list of edges describing DynapcnnLayer connectivity."