| 88 | g_a = g["AA"] |
| 89 | |
| 90 | def check_features(g, bg): |
| 91 | for ntype in bg.srctypes: |
| 92 | for key in g.nodes[ntype].data: |
| 93 | assert F.array_equal( |
| 94 | bg.srcnodes[ntype].data[key], |
| 95 | F.gather_row( |
| 96 | g.nodes[ntype].data[key], |
| 97 | bg.srcnodes[ntype].data[dgl.NID], |
| 98 | ), |
| 99 | ) |
| 100 | for ntype in bg.dsttypes: |
| 101 | for key in g.nodes[ntype].data: |
| 102 | assert F.array_equal( |
| 103 | bg.dstnodes[ntype].data[key], |
| 104 | F.gather_row( |
| 105 | g.nodes[ntype].data[key], |
| 106 | bg.dstnodes[ntype].data[dgl.NID], |
| 107 | ), |
| 108 | ) |
| 109 | for etype in bg.canonical_etypes: |
| 110 | for key in g.edges[etype].data: |
| 111 | assert F.array_equal( |
| 112 | bg.edges[etype].data[key], |
| 113 | F.gather_row( |
| 114 | g.edges[etype].data[key], bg.edges[etype].data[dgl.EID] |
| 115 | ), |
| 116 | ) |
| 117 | |
| 118 | bg = dgl.to_block(g_a) |
| 119 | check(g_a, bg, "A", "AA", None) |