relax the criterion on num_edges for using the "sparse" kernel

src/pyjuice/layer/sum_layer.py

@@ -927,7 +927,7 @@ def _forward_sparse(self, node_mars: torch.Tensor, element_mars: torch.Tensor,
         batch_size = node_mars.size(1)
         BATCH_SIZE_NP2 = triton.next_power_of_2(batch_size)
 
-        assert num_edges <= 16384, "The sparse forward kernel only support nodes with # edges smaller than 16384."
+        # assert num_edges <= 16384, "The sparse forward kernel only support nodes with # edges smaller than 16384."
 
         if triton.cdiv(layer_n_nodes, self.block_size) <= 2048:
             BLOCK_B = max(min(2048 // num_edges, BATCH_SIZE_NP2), 1)
@@ -1083,8 +1083,10 @@ def _backward(self, node_flows: torch.Tensor, element_flows: torch.Tensor,
         elif self.block_size * batch_size < 32:
             # Advantage of block-sparse processing is diminishing
             mode = self.SPARSE
-        else:
+        elif num_edges <= 32768:
             mode = self.BLOCK_SPARSE
+        else:
+            mode = self.SPARSE
 
         if mode == self.BLOCK_SPARSE:
             self._backward_block_sparse(
@@ -2320,8 +2322,7 @@ def _backward_sparse_par_flows(self, node_flows: torch.Tensor, params: torch.Ten
         batch_size = node_mars.size(1)
         BATCH_SIZE_NP2 = triton.next_power_of_2(batch_size)
 
-        assert num_edges <= 16384, "The sparse backward kernel only support nodes with # edges smaller than 16384."
-
+        # assert num_edges <= 16384, "The sparse backward kernel only support nodes with # edges smaller than 16384."
 
         if num_edges <= 1024:
             BLOCK_B = max(min(2048 // num_edges, BATCH_SIZE_NP2), 1)