Merge remote-tracking branch 'origin/main' into normalize

Author: PoorvaGarg

src/pyjuice/layer/sum_layer.py

@@ -2637,6 +2637,7 @@ def _bk_triton_large_sparse_ele_kernel(node_flows, element_flows, node_mars, ele
 
                     elflows_max = tl.max(elflows, axis = 1)
                     eflows = tl.log(tl.sum(tl.exp(elflows - elflows_max[:,None,:]), axis = 1)) + elflows_max
+                    eflows = tl.where((elflows_max == -float("inf")) | (emars == -float("inf")), -float("inf"), eflows)
                 else:
                     if propagation_alg_id == 0:
                         eflows = tl.sum(nflows * epars[:,:,None] * tl.exp(emars[:,None,:] - nmars), axis = 1)
@@ -2741,7 +2742,6 @@ def _backward_sparse_ele_flows(self, node_flows: torch.Tensor, element_flows: to
                 )
 
             else:
-
                 for pid_m_start in range(0, grid[1], 32768):
 
                     pid_m_end = min(pid_m_start + 32768, grid[1])