Add an option to tie input parameters in HCLTs

src/pyjuice/structures/compilation.py

@@ -15,7 +15,8 @@ def BayesianTreeToHiddenRegionGraph(tree: nx.Graph,
                                     InputDist: Type[Distribution], 
                                     dist_params: dict,
                                     num_root_ns: int = 1,
-                                    block_size: Optional[int] = None) -> CircuitNodes:
+                                    block_size: Optional[int] = None,
+                                    tie_input_params: bool = False) -> CircuitNodes:
     """
     Given a Tree Bayesian Network tree T1 (i.e. at most one parents), 
     
@@ -48,6 +49,22 @@ def children(n: int):
     for n in clt.nodes:
         assert len(list(clt.predecessors(n))) <= 1
 
+    # For input parameter tying
+    template_ni = None
+
+    def get_input_ns(v):
+        nonlocal template_ni
+        if tie_input_params:
+            if template_ni is None:
+                ni = inputs(v, num_node_blocks = num_node_blocks, dist = InputDist(**dist_params))
+                template_ni = ni
+            else:
+                ni = template_ni.duplicate(v, tie_params = True)
+        else:
+            ni = inputs(v, num_node_blocks = num_node_blocks, dist = InputDist(**dist_params))
+
+        return ni
+
     # Compile the region graph for the circuit equivalent to T2
     node_seq = list(nx.dfs_postorder_nodes(tree, root))
     var2rnode = dict()
@@ -57,7 +74,7 @@ def children(n: int):
 
             if len(chs) == 0:
                 # Input Region
-                r = inputs(v, num_node_blocks = num_node_blocks, dist = InputDist(**dist_params))
+                r = get_input_ns(v)
                 var2rnode[v] = r
             else:
                 # Inner Region
@@ -66,7 +83,7 @@ def children(n: int):
                 ch_regions = [var2rnode[c] for c in chs]
 
                 # Add x_v to children(z_v)
-                leaf_r = inputs(v, num_node_blocks = num_node_blocks, dist = InputDist(**dist_params))
+                leaf_r = get_input_ns(v)
                 ch_regions.append(leaf_r)
 
                 rp = multiply(*ch_regions)

src/pyjuice/structures/hclt.py

@@ -72,7 +72,8 @@ def HCLT(x: torch.Tensor, num_latents: int,
          block_size: Optional[int] = None,
          input_dist: Optional[Distribution] = None,
          input_node_type: Type[Distribution] = Categorical, 
-         input_node_params: dict = {"num_cats": 256}):
+         input_node_params: dict = {"num_cats": 256},
+         tie_input_params: bool = False):
     """
     Construct Hidden Chow-Liu Trees (https://arxiv.org/pdf/2106.02264.pdf).
 
@@ -99,6 +100,9 @@ def HCLT(x: torch.Tensor, num_latents: int,
 
     :param input_dist: input distribution
     :type input_dist: Distribution
+
+    :param tie_input_params: whether to tie the input parameters
+    :type tie_input_params: bool
     """
 
     if input_dist is not None:
@@ -111,7 +115,7 @@ def HCLT(x: torch.Tensor, num_latents: int,
     root_r = BayesianTreeToHiddenRegionGraph(
         T, root, num_latents, input_node_type, 
         input_node_params, num_root_ns = num_root_ns,
-        block_size = block_size
+        block_size = block_size, tie_input_params = tie_input_params
     )
 
     return root_r