Function test_gin_conv

tests/python/tensorflow/test_nn.py:428–434

(g, idtype, aggregator_type)

@pytest.mark.parametrize("g", get_cases(["homo", "block-bipartite"]))
@pytest.mark.parametrize("aggregator_type", ["mean", "max", "sum"])
def test_gin_conv(g, idtype, aggregator_type):
    g = g.astype(idtype).to(F.ctx())
    ctx = F.ctx()
    gin = nn.GINConv(tf.keras.layers.Dense(12), aggregator_type)
    feat = F.randn((g.number_of_src_nodes(), 5))
    h = gin(g, feat)
    assert h.shape == (g.number_of_dst_nodes(), 12)


@parametrize_idtype