def test_column_subcolumn():
data = F.copy_to(
F.tensor([[1., 1., 1., 1.], [0., 2., 9., 0.], [3., 2., 1., 0.],
[1., 1., 1., 1.], [0., 2., 4., 0.]]), F.ctx())
original = Column(data)
# subcolumn from cpu context
i1 = F.tensor([0, 2, 1, 3], dtype=F.int64)
l1 = original.subcolumn(i1)
assert len(l1) == i1.shape[0]
assert F.array_equal(l1.data, F.gather_row(data, i1))
# next subcolumn from target context
i2 = F.copy_to(F.tensor([0, 2], dtype=F.int64), F.ctx())
l2 = l1.subcolumn(i2)
assert len(l2) == i2.shape[0]
i1i2 = F.copy_to(F.gather_row(i1, F.copy_to(i2, F.context(i1))), F.ctx())
assert F.array_equal(l2.data, F.gather_row(data, i1i2))
# next subcolumn also from target context
i3 = F.copy_to(F.tensor([1], dtype=F.int64), F.ctx())
l3 = l2.subcolumn(i3)
assert len(l3) == i3.shape[0]
i1i2i3 = F.copy_to(F.gather_row(i1i2, F.copy_to(i3, F.context(i1i2))),
F.ctx())
assert F.array_equal(l3.data, F.gather_row(data, i1i2i3))
def test_subframes(parent_idx_device, child_device):
parent_device, idx_device = parent_idx_device
g = dgl.graph((F.tensor([1,2,3], dtype=F.int64), F.tensor([2,3,4], dtype=F.int64)))
print(g.device)
g.ndata['x'] = F.randn((5, 4))
g.edata['a'] = F.randn((3, 6))
idx = F.tensor([1, 2], dtype=F.int64)
if parent_device == 'cuda':
g = g.to(F.cuda())
elif parent_device == 'uva':
if F.backend_name != 'pytorch':
pytest.skip("UVA only supported for PyTorch")
g = g.to(F.cpu())
g.create_formats_()
g.pin_memory_()
elif parent_device == 'cpu':
g = g.to(F.cpu())
idx = F.copy_to(idx, idx_device)
sg = g.sample_neighbors(idx, 2).to(child_device)
assert sg.device == F.context(sg.ndata['x'])
assert sg.device == F.context(sg.edata['a'])
assert sg.device == child_device
if parent_device != 'uva':
sg = g.to(child_device).sample_neighbors(F.copy_to(idx, child_device), 2)
assert sg.device == F.context(sg.ndata['x'])
assert sg.device == F.context(sg.edata['a'])
assert sg.device == child_device
if parent_device == 'uva':
g.unpin_memory_()
def test_async_transferer_to_other():
cpu_ones = F.ones([100,75,25], dtype=F.int32, ctx=F.cpu())
tran = AsyncTransferer(F.ctx())
t = tran.async_copy(cpu_ones, F.ctx())
other_ones = t.wait()
assert F.context(other_ones) == F.ctx()
assert F.array_equal(F.copy_to(other_ones, ctx=F.cpu()), cpu_ones)