fix blk sparse backward with -inf chs

Author: liuanji

src/pyjuice/layer/sum_layer.py

@@ -1688,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,
@@ -1847,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,