Add optimization algorithms and stabilize the backward pass

pyproject.toml

@@ -21,14 +21,26 @@ authors = [
     {name="StarAI", email="[email protected]"},
 ]
 
+[project.optional-dependencies]
+dev = [
+    "pytest",
+    "pytest-xdist",
+    "pytest-skip-slow",
+    "torchvision",
+    "torchtext",
+    "matplotlib"
+]
+
 [options.packages.find]
 where = "src"
 
 [tool.setuptools.dynamic]
 readme = {file = "README.md"}
 
-
 [tool.pytest.ini_options]
 addopts = [
     "--import-mode=importlib",
 ]
+testpaths = [
+    "tests/"
+]

src/pyjuice/layer/input_layer.py

@@ -30,6 +30,8 @@ def __init__(self, nodes: Sequence[InputNodes], cum_nodes: int = 0, pc_num_vars:
                                            in gradient accumulation.
         """
 
+        assert len(nodes) == len(set(nodes)), "Input node list contains duplicates."
+
         nn.Module.__init__(self)
         Layer.__init__(self, nodes, disable_block_size_check = True)
 
@@ -213,7 +215,7 @@ def init_param_flows(self, flows_memory: float = 1.0):
 
     def forward(self, data: torch.Tensor, node_mars: torch.Tensor, params: Optional[Dict] = None,
                 missing_mask: Optional[torch.Tensor] = None, _batch_first: bool = True, 
-                _apply_missing_mask_only: bool = False):
+                _apply_missing_mask_only: bool = False, **kwargs):
         self._used_external_params = (params is not None)
 
         if params is None:
@@ -300,7 +302,8 @@ def forward(self, data: torch.Tensor, node_mars: torch.Tensor, params: Optional[
             raise NotImplementedError("CPU forward fn for input nodes is not implemented.")
 
     def backward(self, data: torch.Tensor, node_flows: torch.Tensor, 
-                 node_mars: torch.Tensor, params: Optional[Dict] = None):
+                 node_mars: torch.Tensor, params: Optional[Dict] = None,
+                 logspace_flows: bool = False, **kwargs):
         """
         data: [num_vars, B]
         node_flows: [num_nodes, B]
@@ -355,6 +358,7 @@ def backward(self, data: torch.Tensor, node_flows: torch.Tensor,
                 node_offset = node_offset, 
                 BLOCK_SIZE = BLOCK_SIZE, 
                 partial_eval = 1 if bk_local_ids is not None else 0,
+                logspace_flows = logspace_flows,
                 num_warps = 8
             )
 
@@ -681,7 +685,7 @@ def _fw_missing_mask_kernel(missing_mask_ptr, node_mars_ptr, vids_ptr, fw_local_
 
     @staticmethod
     def _flows_kernel_template(flow_fn, params_ptr, param_flows_ptr, node_flows_ptr, node_mars_ptr, data_ptr, vids_ptr, s_pids_ptr, s_pfids_ptr,
-                               metadata_ptr, s_mids_ptr, bk_local_ids_ptr, partial_eval: tl.constexpr, layer_num_nodes: tl.constexpr, 
+                               metadata_ptr, s_mids_ptr, bk_local_ids_ptr, partial_eval: tl.constexpr, logspace_flows: tl.constexpr, layer_num_nodes: tl.constexpr, 
                                batch_size: tl.constexpr, num_vars_per_node: tl.constexpr, nv_block_size: tl.constexpr, node_offset: tl.constexpr, 
                                BLOCK_SIZE: tl.constexpr):
         pid = tl.program_id(axis = 0)
@@ -720,6 +724,9 @@ def _flows_kernel_template(flow_fn, params_ptr, param_flows_ptr, node_flows_ptr,
         ns_offsets = (local_offsets + node_offset) * batch_size + batch_offsets
         flows = tl.load(node_flows_ptr + ns_offsets, mask = mask, other = 0)
 
+        if logspace_flows:
+            flows = tl.exp(flows)
+
         flow_fn(local_offsets, ns_offsets, data, flows, node_mars_ptr, params_ptr, param_flows_ptr, s_pids, s_pfids, metadata_ptr, 
                 s_mids_ptr, mask, num_vars_per_node, BLOCK_SIZE)
 

src/pyjuice/layer/layer.py

@@ -7,6 +7,13 @@
 
 
 class Layer():
+
+    propagation_alg_mapping = {
+        "LL": 0,
+        "MPE": 1,
+        "GeneralLL": 2
+    }
+
     def __init__(self, nodes: Sequence[CircuitNodes], disable_block_size_check: bool = False) -> None:
 
         if disable_block_size_check:
@@ -60,3 +67,13 @@ def disable_partial_evaluation(self, forward: bool = True, backward: bool = True
 
     def provided(self, var_name):
         return hasattr(self, var_name) and getattr(self, var_name) is not None
+
+    def _get_propagation_alg_kwargs(self, propagation_alg: str, **kwargs):
+        if propagation_alg == "LL":
+            return {"alpha": 0.0}
+        elif propagation_alg == "MPE":
+            return {"alpha": 0.0}
+        elif propagation_alg == "GeneralLL":
+            return {"alpha": kwargs["alpha"]}
+        else:
+            raise ValueError(f"Unknown propagation algorithm {propagation_alg}.")