| 293 | F._default_context_str == "gpu", reason="GPU pack traces not implemented" |
| 294 | ) |
| 295 | def test_pack_traces(): |
| 296 | traces, types = ( |
| 297 | np.array( |
| 298 | [[0, 1, -1, -1, -1, -1, -1], [0, 1, 1, 3, 0, 0, 0]], dtype="int64" |
| 299 | ), |
| 300 | np.array([0, 0, 1, 0, 0, 1, 0], dtype="int64"), |
| 301 | ) |
| 302 | traces = F.zerocopy_from_numpy(traces) |
| 303 | types = F.zerocopy_from_numpy(types) |
| 304 | result = dgl.sampling.pack_traces(traces, types) |
| 305 | assert F.array_equal( |
| 306 | result[0], F.tensor([0, 1, 0, 1, 1, 3, 0, 0, 0], dtype=F.int64) |
| 307 | ) |
| 308 | assert F.array_equal( |
| 309 | result[1], F.tensor([0, 0, 0, 0, 1, 0, 0, 1, 0], dtype=F.int64) |
| 310 | ) |
| 311 | assert F.array_equal(result[2], F.tensor([2, 7], dtype=F.int64)) |
| 312 | assert F.array_equal(result[3], F.tensor([0, 2], dtype=F.int64)) |
| 313 | |
| 314 | |
| 315 | @pytest.mark.parametrize("use_uva", [True, False]) |