| 3310 | |
| 3311 | @parametrize_idtype |
| 3312 | def test_frame(idtype): |
| 3313 | g = dgl.graph(([0, 1, 2], [1, 2, 3]), idtype=idtype, device=F.ctx()) |
| 3314 | g.ndata["h"] = F.copy_to(F.tensor([0, 1, 2, 3], dtype=idtype), ctx=F.ctx()) |
| 3315 | g.edata["h"] = F.copy_to(F.tensor([0, 1, 2], dtype=idtype), ctx=F.ctx()) |
| 3316 | |
| 3317 | # remove nodes |
| 3318 | sg = dgl.remove_nodes(g, [3]) |
| 3319 | # check for lazy update |
| 3320 | assert F.array_equal(sg._node_frames[0]._columns["h"].storage, g.ndata["h"]) |
| 3321 | assert F.array_equal(sg._edge_frames[0]._columns["h"].storage, g.edata["h"]) |
| 3322 | assert sg.ndata["h"].shape[0] == 3 |
| 3323 | assert sg.edata["h"].shape[0] == 2 |
| 3324 | # update after read |
| 3325 | assert F.array_equal( |
| 3326 | sg._node_frames[0]._columns["h"].storage, |
| 3327 | F.tensor([0, 1, 2], dtype=idtype), |
| 3328 | ) |
| 3329 | assert F.array_equal( |
| 3330 | sg._edge_frames[0]._columns["h"].storage, F.tensor([0, 1], dtype=idtype) |
| 3331 | ) |
| 3332 | |
| 3333 | ng = dgl.add_nodes(sg, 1) |
| 3334 | assert ng.ndata["h"].shape[0] == 4 |
| 3335 | assert F.array_equal( |
| 3336 | ng._node_frames[0]._columns["h"].storage, |
| 3337 | F.tensor([0, 1, 2, 0], dtype=idtype), |
| 3338 | ) |
| 3339 | ng = dgl.add_edges(ng, [3], [1]) |
| 3340 | assert ng.edata["h"].shape[0] == 3 |
| 3341 | assert F.array_equal( |
| 3342 | ng._edge_frames[0]._columns["h"].storage, |
| 3343 | F.tensor([0, 1, 0], dtype=idtype), |
| 3344 | ) |
| 3345 | |
| 3346 | # multi level lazy update |
| 3347 | sg = dgl.remove_nodes(g, [3]) |
| 3348 | assert F.array_equal(sg._node_frames[0]._columns["h"].storage, g.ndata["h"]) |
| 3349 | assert F.array_equal(sg._edge_frames[0]._columns["h"].storage, g.edata["h"]) |
| 3350 | ssg = dgl.remove_nodes(sg, [1]) |
| 3351 | assert F.array_equal( |
| 3352 | ssg._node_frames[0]._columns["h"].storage, g.ndata["h"] |
| 3353 | ) |
| 3354 | assert F.array_equal( |
| 3355 | ssg._edge_frames[0]._columns["h"].storage, g.edata["h"] |
| 3356 | ) |
| 3357 | # ssg is changed |
| 3358 | assert ssg.ndata["h"].shape[0] == 2 |
| 3359 | assert ssg.edata["h"].shape[0] == 0 |
| 3360 | assert F.array_equal( |
| 3361 | ssg._node_frames[0]._columns["h"].storage, |
| 3362 | F.tensor([0, 2], dtype=idtype), |
| 3363 | ) |
| 3364 | # sg still in lazy model |
| 3365 | assert F.array_equal(sg._node_frames[0]._columns["h"].storage, g.ndata["h"]) |
| 3366 | assert F.array_equal(sg._edge_frames[0]._columns["h"].storage, g.edata["h"]) |
| 3367 | |
| 3368 | |
| 3369 | @unittest.skipIf( |