(idtype)
| 147 | |
| 148 | @parametrize_idtype |
| 149 | def test_batch_send_and_recv(idtype): |
| 150 | t1 = tree1(idtype) |
| 151 | t2 = tree2(idtype) |
| 152 | |
| 153 | bg = dgl.batch([t1, t2]) |
| 154 | _mfunc = lambda edges: {"m": edges.src["h"]} |
| 155 | _rfunc = lambda nodes: {"h": F.sum(nodes.mailbox["m"], 1)} |
| 156 | u = [3, 4, 2 + 5, 0 + 5] |
| 157 | v = [1, 1, 4 + 5, 4 + 5] |
| 158 | |
| 159 | bg.send_and_recv((u, v), _mfunc, _rfunc) |
| 160 | |
| 161 | t1, t2 = dgl.unbatch(bg) |
| 162 | assert F.asnumpy(t1.ndata["h"][1]) == 7 |
| 163 | assert F.asnumpy(t2.ndata["h"][4]) == 2 |
| 164 | |
| 165 | |
| 166 | @parametrize_idtype |
nothing calls this directly
no test coverage detected