prevent possibility of missing input slots

Author: TimeDelta

population.py

@@ -137,6 +137,7 @@ def __init__(self, config):
             self.shared_attr_vocab,
             graph_encoder.pin_role_embedding,
             min_pin_nodes=decoder_min_pin_nodes,
+            required_input_count=len(configured_inputs),
         )
         icnn_hidden_dims = getattr(config, "latent_icnn_hidden_dims", (64, 32))
         if isinstance(icnn_hidden_dims, str):
@@ -1569,6 +1570,10 @@ def _assign_pin_role(idx: int, role: str) -> bool:
                 return False
             attrs = dict(normalized_attrs[idx] or {})
             current = _normalize_pin_role(attrs.get("pin_role"))
+            if attrs.get("_pin_role_locked"):
+                if current == role:
+                    return False
+                return False
             if current == role:
                 return False
             attrs["pin_role"] = role
@@ -1580,6 +1585,10 @@ def _assign_pin_slot(idx: int, slot: int | None) -> bool:
                 return False
             attrs = dict(normalized_attrs[idx] or {})
             current = _node_pin_slot(idx)
+            if attrs.get("_pin_slot_locked"):
+                if current == slot:
+                    return False
+                return False
             if slot is None:
                 if "pin_slot_index" in attrs:
                     attrs.pop("pin_slot_index", None)

search_space_compression.py

@@ -845,6 +845,7 @@ def __init__(
         max_edges_per_node: int = 256,
         max_edges_per_graph: int = 4096,
         min_pin_nodes: int | None = None,
+        required_input_count: int | None = None,
     ):
         super().__init__()
         self.shared_attr_vocab = shared_attr_vocab
@@ -885,6 +886,13 @@ def __init__(
                 self.min_pin_nodes = max(0, int(min_pin_nodes))
             except (TypeError, ValueError):
                 self.min_pin_nodes = 0
+        if required_input_count is None:
+            self.required_input_count = 0
+        else:
+            try:
+                self.required_input_count = max(0, int(required_input_count))
+            except (TypeError, ValueError):
+                self.required_input_count = 0
         # Encourage attribute decoder to terminate by progressively biasing the EOS logit.
         self.attr_eos_bias_base = 0.0
         self.attr_eos_bias_slope = 0.1
@@ -1022,6 +1030,8 @@ def make_name_input(token_idx: int) -> torch.Tensor:
                     node_loop_timer = time.perf_counter()
                     node_loop_iters = 0
                     teacher_force_nodes = self.training and (target_node_count is not None)
+                    forced_pin_roles: List[str | None] = []
+                    forced_pin_slots: List[int | None] = []
                     with torch.autograd.profiler.record_function("GraphDecoder.node_loop"):
                         while True:
                             logger.debug(f"Decoding node {t}")
@@ -1072,6 +1082,13 @@ def make_name_input(token_idx: int) -> torch.Tensor:
                             node_types.append(self.type_head(new_node).argmax(dim=-1).cpu().tolist())
 
                             node_index = t
+                            forced_role = None
+                            forced_slot = None
+                            if self.required_input_count and node_index < self.required_input_count:
+                                forced_role = PIN_ROLE_INPUT
+                                forced_slot = node_index
+                            forced_pin_roles.append(forced_role)
+                            forced_pin_slots.append(forced_slot)
                             hidden_edge = hidden_node
                             edge_in = torch.zeros(1, 1, device=device)
                             node_edge_budget = 0
@@ -1163,6 +1180,12 @@ def make_name_input(token_idx: int) -> torch.Tensor:
                 pin_role_vec = None
                 decoded_role = None
                 target_pin_role = None
+                forced_role = None
+                forced_slot = None
+                if node_idx < len(forced_pin_roles):
+                    forced_role = forced_pin_roles[node_idx]
+                if node_idx < len(forced_pin_slots):
+                    forced_slot = forced_pin_slots[node_idx]
                 if self.pin_role_head is not None:
                     pin_role_vec = self.pin_role_head(embedding)
                     decoded_role = self._infer_pin_role(pin_role_vec, pin_role_reference)
@@ -1176,8 +1199,14 @@ def make_name_input(token_idx: int) -> torch.Tensor:
                         pin_role_teacher_tokens += 1
                 if decoded_role is None:
                     decoded_role = target_pin_role
+                locked_role = False
+                if forced_role is not None:
+                    decoded_role = forced_role
+                    locked_role = True
                 if decoded_role is not None:
                     attrs["pin_role"] = decoded_role
+                    if locked_role:
+                        attrs["_pin_role_locked"] = True
                 slot_scalar = self.pin_slot_head(embedding).squeeze()
                 slot_scalar = _sanitize_pin_slot_scalar(slot_scalar)
                 slot_target = None
@@ -1189,13 +1218,19 @@ def make_name_input(token_idx: int) -> torch.Tensor:
                     pin_slot_teacher_loss = pin_slot_teacher_loss + slot_loss
                     pin_slot_teacher_tokens += 1
                 slot_attr = None
-                if slot_scalar is not None:
+                locked_slot = False
+                if forced_slot is not None:
+                    slot_attr = forced_slot
+                    locked_slot = True
+                elif slot_scalar is not None:
                     slot_value = float(slot_scalar.detach().item())
                     slot_attr = max(0, int(round(abs(slot_value)) - 1))
                 elif slot_target is not None:
                     slot_attr = max(0, int(round(abs(float(slot_target))) - 1))
                 if slot_attr is not None:
                     attrs["pin_slot_index"] = slot_attr
+                    if locked_slot:
+                        attrs["_pin_slot_locked"] = True
                 name_hidden = embedding.unsqueeze(0).unsqueeze(0)
                 val_hidden = None
                 t = 0

tests/test_population.py

@@ -744,6 +744,33 @@ def reaches_outputs():
         assert idx in backward
 
 
+def test_repair_preserves_seeded_input_pins():
+    config = make_config()
+    pop = GuidedPopulation(config)
+    required_inputs = len(config.genome_config.input_keys)
+    if required_inputs < 2:
+        pytest.skip("insufficient configured inputs for this test")
+    node_attrs = []
+    for slot in range(required_inputs):
+        node_attrs.append(
+            {
+                "pin_role": "input",
+                "pin_slot_index": slot,
+                "_pin_role_locked": True,
+                "_pin_slot_locked": True,
+            }
+        )
+    node_attrs.append({})
+    graph = _make_empty_graph_dict(len(node_attrs), node_attrs)
+
+    pop._repair_graph_dict(graph)
+
+    attrs = graph["node_attributes"]
+    for slot in range(required_inputs):
+        assert attrs[slot]["pin_role"] == "input"
+        assert attrs[slot]["pin_slot_index"] == slot
+
+
 def test_repair_preserves_predicted_edges_for_visualization():
     config = make_config()
     required_inputs = len(config.genome_config.input_keys)

tests/test_search_space_compression.py

@@ -231,3 +231,29 @@ def test_graph_decoder_enforces_min_pin_nodes():
     graphs = decoder(latent)
     graph = graphs[0]
     assert len(graph["node_attributes"]) >= 3
+
+
+def test_graph_decoder_seeds_input_pin_roles():
+    vocab = SharedAttributeVocab([], embedding_dim=4)
+    decoder = GraphDecoder(
+        num_node_types=3,
+        latent_dim=8,
+        shared_attr_vocab=vocab,
+        hidden_dim=4,
+        min_pin_nodes=4,
+        required_input_count=2,
+    )
+    decoder.eval()
+
+    with torch.no_grad():
+        decoder.stop_head.weight.zero_()
+        decoder.stop_head.bias.fill_(20.0)
+
+    latent = torch.zeros(1, decoder.latent_dim)
+    graphs = decoder(latent)
+    graph = graphs[0]
+    attrs = graph["node_attributes"]
+    assert attrs[0]["pin_role"] == "input"
+    assert attrs[0]["pin_slot_index"] == 0
+    assert attrs[1]["pin_role"] == "input"
+    assert attrs[1]["pin_slot_index"] == 1