Merge branch 'main' of https://github.com/Tractables/pyjuice into normalize

Author: PoorvaGarg

README.md

@@ -240,3 +240,17 @@ sum_ns = juice.summate(prod_ns1, prod_ns2, num_node_blocks = num_nodes // 4, blo
 ```
 
 The above is equivalent to considering the input nodes to be concatenated into a single vector of nodes, and then define the edges correspondingly.
+
+## Citation
+
+If you find PyJuice useful, please consider citing us:
+
+```
+@inproceedings{liu2024scaling,
+  title         = {Scaling Tractable Probabilistic Circuits: A Systems Perspective},
+  author        = {Liu, Anji and Ahmed, Kareem and Van den Broeck, Guy},
+  booktitle     = {Proceedings of the 41th International Conference on Machine Learning (ICML)},
+  month         = {jul},
+  year          = {2024}
+}
+```

src/pyjuice/layer/compilation.py

@@ -268,7 +268,7 @@ def sum_layer_forward_compilation(nodes, fw_partition_max_chs, n_partition_ids,
         use_cuda = False
 
     if use_cuda:
-        device = torch.device("cuda:0")
+        device = torch.device(f"cuda:{torch.cuda.current_device()}")
     else:
         device = torch.device("cpu")
 
@@ -703,7 +703,7 @@ def sum_layer_backward_compilation(nodes, cs2parns, n_partition_ids, n_id_in_par
         use_cuda = False
 
     if use_cuda:
-        device = torch.device("cuda:0")
+        device = torch.device(f"cuda:{torch.cuda.current_device()}")
     else:
         device = torch.device("cpu")
 
@@ -912,7 +912,7 @@ def prod_layer_forward_compilation(nodes, fw_partition_max_chs, n_partition_ids,
         assert block_size == 1
 
     if use_cuda:
-        device = torch.device("cuda:0")
+        device = torch.device(f"cuda:{torch.cuda.current_device()}")
     else:
         device = torch.device("cpu")
 

src/pyjuice/layer/input_layer.py

@@ -499,7 +499,7 @@ def sample(self, samples: torch.Tensor, node_flows: torch.Tensor, missing_mask:
                 nv_block_size = triton.next_power_of_2(self.num_vars_per_node),
                 batch_size = batch_size, 
                 BLOCK_SIZE = BLOCK_SIZE,
-                seed = seed if seed is not None else random.randint(0, 1e8)
+                seed = seed if seed is not None else random.randint(0, int(1e8))
             )
 
         else:

src/pyjuice/layer/sum_layer.py

@@ -531,7 +531,10 @@ def _fw_triton_block_sparse_tlmm_kernel(node_mars, element_mars, params, nids, c
 
                 acc = tl.where(emars_max > acc,
                     tl.log(nmars + tl.exp(acc - emars_max) + 1e-24) + emars_max,
-                    tl.log(tl.exp(emars_max - acc) * nmars + 1.0) + acc
+                    tl.where(acc != -float("inf"),
+                        tl.log(tl.exp(emars_max - acc) * nmars + 1.0) + acc,
+                        -float("inf")
+                    )
                 )
 
             # Increment `epars_ptr`
@@ -647,7 +650,10 @@ def _fw_triton_block_sparse_csmm1_kernel(node_mars, element_mars, params, nids,
 
                 acc = tl.where(emars_max > acc,
                     tl.log(nmars + tl.exp(acc - emars_max) + 1e-24) + emars_max,
-                    tl.log(tl.exp(emars_max - acc) * nmars + 1.0) + acc
+                    tl.where(acc != -float("inf"),
+                        tl.log(tl.exp(emars_max - acc) * nmars + 1.0) + acc,
+                        -float("inf")
+                    )
                 )
 
             # Increment `epars_ptr`
@@ -757,7 +763,10 @@ def _fw_triton_block_sparse_csmm2_kernel(node_mars, element_mars, params, nids,
 
                 acc = tl.where(emars_max[None,:] > acc,
                     tl.log(nmars + tl.exp(acc - emars_max[None,:]) + 1e-24) + emars_max[None,:],
-                    tl.log(tl.exp(emars_max[None,:] - acc) * nmars + 1.0) + acc
+                    tl.where(acc != -float("inf"), 
+                        tl.log(tl.exp(emars_max[None,:] - acc) * nmars + 1.0) + acc,
+                        -float("inf")
+                    )
                 )
 
             # Increment `epars_ptr`
@@ -1679,7 +1688,7 @@ def _bk_triton_block_sparse_ele_kernel(node_flows, element_flows, node_mars, ele
 
                 if logspace_flows:
                     partial_flows_max = emars + log_n_fdm_max
-                    acc = tl.where(log_n_fdm_max == -float("inf"),
+                    acc = tl.where(partial_flows_max == -float("inf"),
                         acc,
                         tl.where(partial_flows_max > acc,
                             tl.log(partial_flows + tl.exp(acc - partial_flows_max) + 1e-24) + partial_flows_max,
@@ -1838,7 +1847,7 @@ def _bk_triton_block_sparse_ele_csmm2_kernel(node_flows, element_flows, node_mar
 
                 if logspace_flows:
                     partial_flows_max = emars + log_n_fdm_max[None,:]
-                    acc = tl.where(log_n_fdm_max[None,:] == -float("inf"),
+                    acc = tl.where(partial_flows_max == -float("inf"),
                         acc,
                         tl.where(partial_flows_max > acc,
                             tl.log(partial_flows + tl.exp(acc - partial_flows_max) + 1e-24) + partial_flows_max,