Composition: fix incorrect target parsing due to input order (#3208)

Fixes #3203
PrincetonUniversity · Feb 22, 2025 · 6dbe8c7 · 6dbe8c7
1 parent a2133f8
commit 6dbe8c7
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 1 deletion.
diff --git a/psyneulink/core/compositions/composition.py b/psyneulink/core/compositions/composition.py
@@ -10058,7 +10058,7 @@ def _infer_target_nodes(self, targets: dict, execution_mode):
 
         if execution_mode is pnlvm.ExecutionMode.PyTorch:
             # Reassign target inputs from output Nodes to target mechanisms constructed for PyTorch execution
-            return {target: value for target, value in zip(self.targets_from_outputs_map.keys(), targets.values())}
+            return {self.outputs_to_targets_map[target]: value for target, value in targets.items()}
 
         ret = {}
         for node, values in targets.items():

diff --git a/tests/composition/test_composition.py b/tests/composition/test_composition.py
@@ -6600,6 +6600,52 @@ def test_get_input_format(self, form, use_labels, show_nested, num_trials, expec
                 num_trials = 2
             len(ocomp.results)==num_trials
 
+    @pytest.mark.pytorch
+    @pytest.mark.parametrize(
+        'merge_targets_1', [True, False], ids=['merge_targets_1', 'separate_targets_1'],
+    )
+    @pytest.mark.parametrize(
+        'merge_targets_2', [True, False], ids=['merge_targets_2', 'separate_targets_2'],
+    )
+    def test_targets_order(self, merge_targets_1, merge_targets_2):
+        in_1 = pnl.TransferMechanism(name='in_1', input_shapes=1)
+        in_2 = pnl.TransferMechanism(name='in_2', input_shapes=2)
+        out_1 = pnl.TransferMechanism(name='out_1', input_shapes=1)
+        out_2 = pnl.TransferMechanism(name='out_2', input_shapes=2)
+
+        comp = pnl.AutodiffComposition(
+            name='comp',
+            pathways=[
+                [in_1, out_1],
+                [in_1, out_2],
+                [in_2, out_1],
+                [in_2, out_2],
+            ],
+        )
+
+        in_1_train = [0, 0, 0, 1, 1, 1]
+        in_2_train = [[0, 0], [0, 0], [0, 0], [1, 1], [1, 1], [1, 1]]
+        out_1_train = [0, 0, 0, 1, 1, 1]
+        out_2_train = [[0, 0], [0, 0], [0, 0], [1, 1], [1, 1], [1, 1]]
+
+        inp_1 = {
+            'inputs': {in_1: in_1_train, in_2: in_2_train},
+            'targets': {out_1: out_1_train, out_2: out_2_train},
+        }
+        inp_2 = {
+            'inputs': {in_1: in_1_train, in_2: in_2_train},
+            'targets': {out_2: out_2_train, out_1: out_1_train},
+        }
+        if merge_targets_1:
+            inp_1 = {'inputs': inp_1}
+        if merge_targets_2:
+            inp_2 = {'inputs': inp_2}
+
+        res_1 = comp.learn(**inp_1, execution_mode=pnl.ExecutionMode.PyTorch, context=1)
+        res_2 = comp.learn(**inp_2, execution_mode=pnl.ExecutionMode.PyTorch, context=2)
+
+        np.testing.assert_array_equal(res_1, res_2)
+
 
 class TestProperties: