chore: import upstream snapshot with attribution

2026-07-13 13:35:51 +08:00
commit c36a561cd8
2172 changed files with 455595 additions and 0 deletions
@@ -0,0 +1,209 @@
+import random
+
+import backend as F
+
+import dgl
+import numpy as np
+import pytest
+import torch
+from utils import parametrize_idtype
+
+random.seed(42)
+np.random.seed(42)
+dgl.seed(42)
+torch.random.manual_seed(42)
+
+
+@parametrize_idtype
+@pytest.mark.parametrize("feat_size", [(5,), ()])
+def test_copy_u(idtype, feat_size):
+    ctx = F.ctx()
+    g = dgl.rand_graph(30, 100)
+    g = g.astype(idtype).to(ctx)
+    x = torch.randn(
+        (g.num_nodes(),) + feat_size, requires_grad=True, device=ctx
+    )
+
+    y = dgl.copy_u(g, x)
+    y.sum().backward()
+    x_grad = x.grad
+
+    x.grad.zero_()
+    u, v = g.edges()
+    y_true = x[u.long()]
+    y_true.sum().backward()
+    x_grad_true = x.grad
+
+    assert torch.allclose(y, y_true)
+    assert torch.allclose(x_grad, x_grad_true)
+
+
+@parametrize_idtype
+@pytest.mark.parametrize("feat_size", [(5,), ()])
+def test_copy_u_hetero(idtype, feat_size):
+    ctx = F.ctx()
+    hg = dgl.heterograph(
+        {
+            ("user", "follow", "user"): ([0, 1, 2], [2, 3, 4]),
+            ("user", "like", "movie"): ([3, 3, 1, 2], [0, 0, 1, 1]),
+        }
+    )
+
+    hg = hg.astype(idtype).to(ctx)
+    x = torch.randn(
+        (hg.num_nodes("user"),) + feat_size, requires_grad=True, device=ctx
+    )
+
+    y = dgl.copy_u(hg, x, etype="like")
+    y.sum().backward()
+    x_grad = x.grad
+
+    x.grad.zero_()
+    u, v = hg.edges(etype="like")
+    y_true = x[u.long()]
+    y_true.sum().backward()
+    x_grad_true = x.grad
+
+    assert torch.allclose(y, y_true)
+    assert torch.allclose(x_grad, x_grad_true)
+
+
+@parametrize_idtype
+@pytest.mark.parametrize("feat_size", [(5,), ()])
+def test_copy_v(idtype, feat_size):
+    ctx = F.ctx()
+    g = dgl.rand_graph(30, 100)
+    g = g.astype(idtype).to(ctx)
+    x = torch.randn(
+        (g.num_nodes(),) + feat_size, requires_grad=True, device=ctx
+    )
+
+    y = dgl.copy_v(g, x)
+    y.sum().backward()
+    x_grad = x.grad
+
+    x.grad.zero_()
+    u, v = g.edges()
+    y_true = x[v.long()]
+    y_true.sum().backward()
+    x_grad_true = x.grad
+
+    assert torch.allclose(y, y_true)
+    assert torch.allclose(x_grad, x_grad_true)
+
+
+@parametrize_idtype
+@pytest.mark.parametrize("feat_size", [(5,), ()])
+def test_copy_v_hetero(idtype, feat_size):
+    ctx = F.ctx()
+    hg = dgl.heterograph(
+        {
+            ("user", "follow", "user"): ([0, 1, 2], [2, 3, 4]),
+            ("user", "like", "movie"): ([3, 3, 1, 2], [0, 0, 1, 1]),
+        }
+    )
+
+    hg = hg.astype(idtype).to(ctx)
+    x = torch.randn(
+        (hg.num_nodes("movie"),) + feat_size, requires_grad=True, device=ctx
+    )
+
+    y = dgl.copy_v(hg, x, etype="like")
+    y.sum().backward()
+    x_grad = x.grad
+
+    x.grad.zero_()
+    u, v = hg.edges(etype="like")
+    y_true = x[v.long()]
+    y_true.sum().backward()
+    x_grad_true = x.grad
+
+    assert torch.allclose(y, y_true)
+    assert torch.allclose(x_grad, x_grad_true)
+
+
+binary_arg_sizes = [
+    ((5,), (5,)),
+    ((5,), ()),
+    ((), (5,)),
+    ((1, 3, 3), (4, 1, 3)),
+    ((3, 3), (4, 1, 3)),
+    ((4, 1, 3), (3, 3)),
+]
+
+dot_arg_sizes = [
+    ((5,), (5,)),
+    ((1, 3, 3), (4, 1, 3)),
+    ((3, 3), (4, 1, 3)),
+    ((4, 1, 3), (3, 3)),
+]
+
+ops = ["add", "sub", "mul", "div"]
+
+
+def pad_shape(x, y, x_size, y_size):
+    xy_size = torch.broadcast_shapes(x_size, y_size)
+    new_x_size = (1,) * (len(xy_size) - len(x_size)) + x_size
+    new_y_size = (1,) * (len(xy_size) - len(y_size)) + y_size
+    new_x = x.view(-1, *new_x_size)
+    new_y = y.view(-1, *new_y_size)
+    return new_x, new_y
+
+
+@parametrize_idtype
+@pytest.mark.parametrize("op", ops)
+@pytest.mark.parametrize("x_size,y_size", binary_arg_sizes)
+def test_u_op_v(idtype, op, x_size, y_size):
+    ctx = F.ctx()
+    g = dgl.rand_graph(30, 100)
+    g = g.astype(idtype).to(ctx)
+    x = torch.randn((g.num_nodes(),) + x_size, requires_grad=True, device=ctx)
+    y = torch.randn((g.num_nodes(),) + y_size, requires_grad=True, device=ctx)
+
+    f_dgl = getattr(dgl, f"u_{op}_v")
+    z = f_dgl(g, x, y)
+    z.sum().backward()
+    x_grad = x.grad
+    y_grad = y.grad
+
+    x_grad.zero_()
+    y_grad.zero_()
+    u, v = g.edges()
+    f_torch = getattr(torch, op)
+    x_u, y_v = pad_shape(x[u.long()], y[v.long()], x_size, y_size)
+    z_true = f_torch(x_u, y_v)
+    z_true.sum().backward()
+    x_grad_true = x.grad
+    y_grad_true = y.grad
+
+    assert torch.allclose(z, z_true)
+    assert torch.allclose(x_grad, x_grad_true)
+    assert torch.allclose(y_grad, y_grad_true)
+
+
+@parametrize_idtype
+@pytest.mark.parametrize("x_size,y_size", dot_arg_sizes)
+def test_u_dot_v(idtype, x_size, y_size):
+    ctx = F.ctx()
+    g = dgl.rand_graph(30, 100)
+    g = g.astype(idtype).to(ctx)
+    x = torch.randn((g.num_nodes(),) + x_size, requires_grad=True, device=ctx)
+    y = torch.randn((g.num_nodes(),) + y_size, requires_grad=True, device=ctx)
+
+    z = dgl.u_dot_v(g, x, y)
+    z.sum().backward()
+    x_grad = x.grad
+    y_grad = y.grad
+
+    x_grad.zero_()
+    y_grad.zero_()
+    u, v = g.edges()
+    x_u, y_v = pad_shape(x[u.long()], y[v.long()], x_size, y_size)
+    z_true = (x_u * y_v).sum(-1).unsqueeze(-1)
+    z_true.sum().backward()
+    x_grad_true = x.grad
+    y_grad_true = y.grad
+
+    assert torch.allclose(z, z_true, atol=1e-4, rtol=1e-4)
+    assert torch.allclose(x_grad, x_grad_true)
+    assert torch.allclose(y_grad, y_grad_true)