fix triton bug in sampling

Author: liuanji

src/pyjuice/queries/sample.py

@@ -197,6 +197,9 @@ def count_prod_nch_kernel(nids, cids, element_samples, ind_ch_count, ind_nids, i
     tl.store(ind_nid_offs + offs_sample, local_nid_offs, mask = mask_sample)
     tl.store(ind_mask + offs_sample, partition_id, mask = mask_sample)
 
+    # Handle triton bug.. (otherwise `local_nids` will be wrong)
+    local_nids = tl.load(ind_nids + offs_sample, mask = mask_sample, other = 0)
+
     # Offset for children
     offs_child = tl.arange(0, BLOCK_C)
     mask_child = offs_child < num_edges

tests/queries/cond_sample_test.py

@@ -129,4 +129,5 @@ def test_cond_sample():
 
 
 if __name__ == "__main__":
+    torch.manual_seed(2389)
     test_cond_sample()