add num_ch_nodes for sum nodes

Tractables · Mar 14, 2024 · e298620 · e298620
1 parent ab66287
commit e298620
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/pyjuice/nodes/sum_nodes.py b/src/pyjuice/nodes/sum_nodes.py
@@ -71,6 +71,10 @@ def num_edges(self):
         """
         return self.edge_ids.size(1) * self.block_size * self.ch_block_size
 
+    @property
+    def num_ch_nodes(self):
+        return self.num_ch_node_blocks * self.ch_block_size
+
     def duplicate(self, *args, tie_params: bool = False) -> SumNodes:
         """
         Create a duplication of the current node with the same specification (i.e., number of nodes, block size).