perf: speed up graph parent generation and idempotent checks via size pre-filtering

src/gflownet/algo/trajectory_balance.py

@@ -276,14 +276,17 @@ def get_idempotent_actions(self, g: Graph, gd: gd.Data, gp: Graph, action: Graph
         lmask = getattr(gd, action.action.mask_name)
         nz = lmask.nonzero()  # Legal actions are those with a nonzero mask value
         actions = [iaction if return_aidx else action]
+        gp_n = len(gp)
+        gp_e = gp.number_of_edges()
         for i in nz:
             aidx = ActionIndex(action_type=iaction[0], row_idx=i[0].item(), col_idx=i[1].item())
             if aidx == iaction:
                 continue
             ga = self.ctx.ActionIndex_to_GraphAction(gd, aidx, fwd=not action.action.is_backward)
             child = self.env.step(g, ga)
-            if nx.algorithms.is_isomorphic(child, gp, lambda a, b: a == b, lambda a, b: a == b):
-                actions.append(aidx if return_aidx else ga)
+            if len(child) == gp_n and child.number_of_edges() == gp_e:
+                if nx.algorithms.is_isomorphic(child, gp, lambda a, b: a == b, lambda a, b: a == b):
+                    actions.append(aidx if return_aidx else ga)
         return actions
 
     def construct_batch(self, trajs, cond_info, log_rewards):
@@ -774,4 +777,4 @@ def subtb_cum(self, P_F, P_B, F, R, traj_lengths):
             fr = torch.cat([F[s_idx:e_idx], torch.tensor([R[ep]], device=F.device)])
             p = pdiff[s_idx:e_idx]
             total_loss[ep] = self._loss(subTB(fr, p)).sum() / (n * n + n) * 2
-        return total_loss
+        return total_loss

src/gflownet/envs/graph_building_env.py

@@ -249,7 +249,7 @@ def parents(self, g: Graph):
         parents: List[Pair(GraphAction, Graph)]
             The list of parent-action pairs that lead to `g`.
         """
-        parents: List[Tuple[GraphAction, Graph]] = []
+        parents: List[Tuple[GraphAction, Graph, int, int]] = []
         # Count node degrees
         degree: Dict[int, int] = defaultdict(int)
         for a, b in g.edges:
@@ -259,11 +259,14 @@ def parents(self, g: Graph):
         def add_parent(a, new_g):
             # Only add parent if the proposed parent `new_g` is not isomorphic
             # to already identified parents
-            for ap, gp in parents:
+            n_nodes = len(new_g)
+            n_edges = new_g.number_of_edges()
+            for ap, gp, gp_n, gp_e in parents:
                 # Here we are relying on the dict equality operator for nodes and edges
-                if is_isomorphic(new_g, gp, lambda a, b: a == b, lambda a, b: a == b):
-                    return
-            parents.append((a, new_g))
+                if n_nodes == gp_n and n_edges == gp_e:
+                    if is_isomorphic(new_g, gp, lambda a, b: a == b, lambda a, b: a == b):
+                        return
+            parents.append((a, new_g, n_nodes, n_edges))
 
         for a, b in g.edges:
             if degree[a] > 1 and degree[b] > 1 and len(g.edges[(a, b)]) == 0:
@@ -305,7 +308,7 @@ def add_parent(a, new_g):
                     GraphAction(GraphActionType.SetNodeAttr, source=i, attr=k, value=g.nodes[i][k]),
                     graph_without_node_attr(g, i, k),
                 )
-        return parents
+        return [(ap, gp) for ap, gp, _, _ in parents]
 
     def count_backward_transitions(self, g: Graph, check_idempotent: bool = False):
         """Counts the number of parents of g (by default, without checking for isomorphisms)"""
@@ -1024,4 +1027,4 @@ def traj_log_n(self, traj):
 def action_type_to_mask(t: GraphActionType, gbatch: gd.Batch, assert_mask_exists: bool = False):
     if assert_mask_exists:
         assert hasattr(gbatch, t.mask_name), f"Mask {t.mask_name} not found in graph data"
-    return getattr(gbatch, t.mask_name) if hasattr(gbatch, t.mask_name) else torch.ones((1, 1), device=gbatch.x.device)
+    return getattr(gbatch, t.mask_name) if hasattr(gbatch, t.mask_name) else torch.ones((1, 1), device=gbatch.x.device)