normalizing a pc

Author: PoorvaGarg

src/pyjuice/model/tensorcircuit.py

@@ -101,7 +101,8 @@ def __init__(self, root_ns: CircuitNodes, layer_sparsity_tol: float = 0.5,
                  force_gpu_compilation: bool = False,
                  max_tied_ns_per_parflow_block: int = 8,
                  device: Optional[Union[int,torch.device]] = None,
-                 verbose: bool = True) -> None:
+                 verbose: bool = True,
+                 normalize: bool = True) -> None:
 
         super(TensorCircuit, self).__init__()
 
@@ -123,7 +124,8 @@ def __init__(self, root_ns: CircuitNodes, layer_sparsity_tol: float = 0.5,
             force_gpu_compilation = force_gpu_compilation,
             max_tied_ns_per_parflow_block = max_tied_ns_per_parflow_block,
             device = device,
-            verbose = verbose
+            verbose = verbose,
+            normalize = normalize
         )
 
         # Hyperparameters for backward pass
@@ -792,7 +794,7 @@ def _get_num_vars(self, ns: CircuitNodes):
 
     def _init_layers(self, layer_sparsity_tol: Optional[float] = None, max_num_partitions: Optional[int] = None,
                      disable_gpu_compilation: bool = False, force_gpu_compilation: bool = False, 
-                     max_tied_ns_per_parflow_block: int = 8, verbose: bool = True, device: Optional[Union[str,torch.device]] = None):
+                     max_tied_ns_per_parflow_block: int = 8, verbose: bool = True, device: Optional[Union[str,torch.device]] = None, normalize: bool = True):
 
         if hasattr(self, "input_layer_group") or hasattr(self, "inner_layer_groups"):
             raise ValueError("Attempting to initialize a TensorCircuit for the second time. " + \
@@ -952,9 +954,9 @@ def _init_layers(self, layer_sparsity_tol: Optional[float] = None, max_num_parti
         self._root_node_range = (self.num_nodes - self.num_root_nodes, self.num_nodes)
 
         # Initialize parameters
-        self._init_parameters()
+        self._init_parameters(normalize = normalize)
 
-    def _init_parameters(self, perturbation: float = 4.0, pseudocount: float = 0.0):
+    def _init_parameters(self, perturbation: float = 4.0, pseudocount: float = 0.0, normalize: bool = True):
         for ns in self.root_ns:
             if not ns.is_tied() and (ns.is_sum() or ns.is_input()) and not ns.has_params():
                 ns.init_parameters(perturbation = perturbation, recursive = False)
@@ -967,7 +969,8 @@ def _init_parameters(self, perturbation: float = 4.0, pseudocount: float = 0.0):
             if ns.is_sum() and not ns.is_tied() and ns.has_params():
                 ns.gather_parameters(params)
 
-        self._normalize_parameters(params, pseudocount = pseudocount)
+        if normalize:
+            self._normalize_parameters(params, pseudocount = pseudocount)
         self.params = nn.Parameter(params)
 
         # Due to the custom inplace backward pass implementation, we do not track 
@@ -982,6 +985,24 @@ def _normalize_parameters(self, params, pseudocount: float = 0.0):
         if params is not None:
             normalize_parameters(params, self.par_update_kwargs, pseudocount)
 
+    def normalize(self, perturbation: float = 0.0, pseudocount: float = 0.0):
+        params = torch.exp(torch.rand([self.num_sum_params]) * -perturbation)
+        params = params.to(self.device)
+        params[:self.num_dummy_params] = 0.0
+
+        # Copy initial parameters if provided
+        for ns in self.root_ns:
+            if ns.is_sum() and not ns.is_tied() and ns.has_params():
+                ns.gather_parameters(params)
+
+        self._normalize_parameters(params, pseudocount = pseudocount)
+        self.params = nn.Parameter(params)
+
+        # Due to the custom inplace backward pass implementation, we do not track 
+        # gradient of PC parameters by PyTorch.
+        self.params.requires_grad = False
+
+
     def _create_node_layers(self):
         depth2nodes = dict()
         nodes2depth = dict()

src/pyjuice/nodes/distributions/literal.py

@@ -74,5 +74,11 @@ def em_fn(local_offsets, params_ptr, param_flows_ptr, s_pids, s_pfids, metadata_
               step_size, pseudocount, BLOCK_SIZE):
         pass
 
+    @staticmethod
+    def partition_fn(local_offsets, params_ptr, s_pids, metadata_ptr, s_mids_ptr, mask, BLOCK_SIZE, TILE_SIZE_K):
+        mars = tl.zeros([BLOCK_SIZE], dtype = tl.float32)
+        return mars
+
+
     def _get_constructor(self):
         return Literal, {"lit": self.lit}