2 files changed, 6 insertions, 1 deletions

diff --git a/pyecsca/sca/re/tree.py b/pyecsca/sca/re/tree.py
index e73a00a..8a05987 100644
--- a/pyecsca/sca/re/tree.py
+++ b/pyecsca/sca/re/tree.py
@@ -133,7 +133,7 @@ class Map:
         has_na = False
         for io_map in mapping.values():
             new = set(io_map.keys())
-            if new != inputs:
+            if inputs and new != inputs:
                 # Map of some cfg doesn't have some inputs, we will fill in None.
                 has_na = True
             inputs.update(new)
diff --git a/test/sca/test_tree.py b/test/sca/test_tree.py
index 03cb041..897ffd5 100644
--- a/test/sca/test_tree.py
+++ b/test/sca/test_tree.py
@@ -22,6 +22,11 @@ def test_map():
     assert dmap.domain == [1, 2]
     assert dmap.codomain == {5, 3, 7, None}
 
+    io_map_full = {"a": {1: 5, 2: 7}, "b": {1: 3, 2: 11}}
+    dmap = Map.from_io_maps(cfgs, io_map_full)
+    assert dmap.domain == [1, 2]
+    assert dmap.codomain == {5, 3, 7, 11}
+
 
 def test_map_merge():
     cfgs = {"a", "b"}