refactor: graph search space cleanup

neps/__init__.py

@@ -7,8 +7,6 @@
     FloatParameter,
     FunctionParameter,
     GraphGrammar,
-    GraphGrammarCell,
-    GraphGrammarRepetitive,
     IntegerParameter,
 )
 from neps.status.status import get_summary_dict, status

neps/search_spaces/__init__.py

@@ -1,8 +1,6 @@
 from neps.search_spaces.architecture.api import ArchitectureParameter, FunctionParameter
 from neps.search_spaces.architecture.graph_grammar import (
     GraphGrammar,
-    GraphGrammarCell,
-    GraphGrammarRepetitive,
     GraphParameter,
 )
 from neps.search_spaces.hyperparameters import (

neps/search_spaces/architecture/api.py

@@ -1,18 +1,20 @@
 from __future__ import annotations
 
 import inspect
-from typing import Callable
+from typing import TYPE_CHECKING, Callable
 
 import networkx as nx
-from torch import nn
 
 from .cfg import Grammar
 from .cfg_variants.constrained_cfg import ConstrainedGrammar
-from .graph_grammar import GraphGrammar, GraphGrammarMultipleRepetitive
+from .graph_grammar import GraphGrammar
+
+if TYPE_CHECKING:
+    from torch import nn
 
 
 def _dict_structure_to_str(
-    structure: dict, primitives: dict, repetitive_mapping: dict = None
+    structure: dict, primitives: dict, repetitive_mapping: dict | None = None
 ) -> str:
     def _save_replace(string: str, __old: str, __new: str):
         while string.count(__old) > 0:
@@ -25,39 +27,36 @@ def _save_replace(string: str, __old: str, __new: str):
     grammar = grammar.replace("(", " ")
     grammar = grammar.replace(")", "")
     grammar = grammar.replace(",", "")
-    for primitive in primitives.keys():
+    for primitive in primitives:
         grammar = _save_replace(grammar, f" {primitive} ", f' "{primitive}" ')
         grammar = _save_replace(grammar, f" {primitive}\n", f' "{primitive}"\n')
     if repetitive_mapping is not None:
-        for placeholder in repetitive_mapping.keys():
+        for placeholder in repetitive_mapping:
             grammar = _save_replace(grammar, f" {placeholder} ", f' "{placeholder}" ')
             grammar = _save_replace(grammar, f" {placeholder}\n", f' "{placeholder}"\n')
     return grammar
 
 
 def _build(graph, set_recursive_attribute):
-    in_node = [n for n in graph.nodes if graph.in_degree(n) == 0][0]
+    in_node = next(n for n in graph.nodes if graph.in_degree(n) == 0)
     for n in nx.topological_sort(graph):
         for pred in graph.predecessors(n):
             e = (pred, n)
             op_name = graph.edges[e]["op_name"]
             if pred == in_node:
                 predecessor_values = None
             else:
-                pred_pred = list(graph.predecessors(pred))[0]
+                pred_pred = next(iter(graph.predecessors(pred)))
                 predecessor_values = graph.edges[(pred_pred, pred)]
             graph.edges[e].update(set_recursive_attribute(op_name, predecessor_values))
 
 
 def ArchitectureParameter(**kwargs):
-    """Factory function"""
-
+    """Factory function."""
     if "structure" not in kwargs:
         raise ValueError("Factory function requires structure")
     if not isinstance(kwargs["structure"], list) or len(kwargs["structure"]) == 1:
         base = GraphGrammar
-    else:
-        base = GraphGrammarMultipleRepetitive
 
     class _FunctionParameter(base):
         def __init__(
@@ -89,9 +88,9 @@ def __init__(
                     _dict_structure_to_str(
                         st,
                         primitives,
-                        repetitive_mapping=kwargs["terminal_to_sublanguage_map"]
-                        if "terminal_to_sublanguage_map" in kwargs
-                        else None,
+                        repetitive_mapping=kwargs.get(
+                            "terminal_to_sublanguage_map", None
+                        ),
                     )
                     if isinstance(st, dict)
                     else st
@@ -144,9 +143,7 @@ def to_pytorch(self) -> nn.Module:
                 self.prune_graph()
 
                 if self._set_recursive_attribute:
-                    m = _build(
-                        self, self._set_recursive_attribute
-                    )
+                    m = _build(self, self._set_recursive_attribute)
 
                 if m is not None:
                     return m