wehub-resource-sync
27 lines
715 B
Python
27 lines
715 B
Python
import io
|
|
|
|
import backend as F
|
|
|
|
import dgl.nn.pytorch as nn
|
|
import pytest
|
|
from utils import parametrize_idtype
|
|
from utils.graph_cases import get_cases
|
|
|
|
tmp_buffer = io.BytesIO()
|
|
|
|
|
|
@parametrize_idtype
|
|
@pytest.mark.parametrize("g", get_cases(["homo"], exclude=["zero-degree"]))
|
|
def test_gatedgcn_conv(g, idtype):
|
|
ctx = F.ctx()
|
|
g = g.astype(idtype).to(ctx)
|
|
gatedgcnconv = nn.GatedGCNConv(10, 10, 5)
|
|
feat = F.randn((g.num_nodes(), 10))
|
|
efeat = F.randn((g.num_edges(), 10))
|
|
gatedgcnconv = gatedgcnconv.to(ctx)
|
|
|
|
h, edge_h = gatedgcnconv(g, feat, efeat)
|
|
# current we only do shape check
|
|
assert h.shape == (g.number_of_dst_nodes(), 5)
|
|
assert edge_h.shape == (g.number_of_edges(), 5)
|