rescale missing flows for input nodes with unnormalized parameters

Author: liuanji

src/pyjuice/layer/input_layer.py

@@ -609,12 +609,17 @@ def eval_partition_fn(self, node_mars: torch.Tensor, params: Optional[Dict] = No
         else:
             raise NotImplementedError("CPU minibatch partition fn for input nodes is not implemented.")
 
-    def add_missing_flows(self, node_flows: torch.Tensor, logspace_flows: bool = False, scale: float = 1.0):
+    def add_missing_flows(self, node_flows: torch.Tensor, node_mars: Optional[torch.Tensor] = None, logspace_flows: bool = False, 
+                          scale: float = 1.0, pc_is_normalized: bool = True):
         """
         Add missing flows specified by `node_flows` to the input node parameters.
         node_flows: [num_nodes]
         """
 
+        if not pc_is_normalized:
+            sid, eid = self._output_ind_range
+            node_flows[sid:eid] /= torch.exp(node_mars[sid:eid])
+
         if "cuda" in self.device.type:
             node_offset = self._output_ind_range[0]
             layer_num_nodes = self._output_ind_range[1] - self._output_ind_range[0]