Function test_gat_conv

tests/python/mxnet/test_nn.py:191–206

(g, idtype, out_dim, num_heads)

@pytest.mark.parametrize("out_dim", [1, 20])
@pytest.mark.parametrize("num_heads", [1, 5])
def test_gat_conv(g, idtype, out_dim, num_heads):
    g = g.astype(idtype).to(F.ctx())
    ctx = F.ctx()
    gat = nn.GATConv(10, out_dim, num_heads)  # n_heads = 5
    gat.initialize(ctx=ctx)
    print(gat)
    feat = F.randn((g.number_of_src_nodes(), 10))
    h = gat(g, feat)
    assert h.shape == (g.number_of_dst_nodes(), num_heads, out_dim)
    _, a = gat(g, feat, True)
    assert a.shape == (g.num_edges(), num_heads, 1)

    # test residual connection
    gat = nn.GATConv(10, out_dim, num_heads, residual=True)
    gat.initialize(ctx=ctx)
    h = gat(g, feat)


@parametrize_idtype