def test_set_trans():
ctx = F.ctx()
g = dgl.DGLGraph(nx.path_graph(15))
st_enc_0 = nn.SetTransformerEncoder(50, 5, 10, 100, 2, 'sab')
st_enc_1 = nn.SetTransformerEncoder(50, 5, 10, 100, 2, 'isab', 3)
st_dec = nn.SetTransformerDecoder(50, 5, 10, 100, 2, 4)
st_enc_0 = st_enc_0.to(ctx)
st_enc_1 = st_enc_1.to(ctx)
st_dec = st_dec.to(ctx)
print(st_enc_0, st_enc_1, st_dec)
# test#1: basic
h0 = F.randn((g.number_of_nodes(), 50))
h1 = st_enc_0(g, h0)
assert h1.shape == h0.shape
h1 = st_enc_1(g, h0)
assert h1.shape == h0.shape
h2 = st_dec(g, h1)
assert h2.shape[0] == 1 and h2.shape[1] == 200 and h2.dim() == 2
# test#2: batched graph
g1 = dgl.DGLGraph(nx.path_graph(5))
g2 = dgl.DGLGraph(nx.path_graph(10))
bg = dgl.batch([g, g1, g2])
h0 = F.randn((bg.number_of_nodes(), 50))
h1 = st_enc_0(bg, h0)
assert h1.shape == h0.shape
h1 = st_enc_1(bg, h0)
assert h1.shape == h0.shape
h2 = st_dec(bg, h1)
assert h2.shape[0] == 3 and h2.shape[1] == 200 and h2.dim() == 2
def test_set_trans():
g = dgl.DGLGraph(nx.path_graph(15))
st_enc_0 = nn.SetTransformerEncoder(50, 5, 10, 100, 2, 'sab')
st_enc_1 = nn.SetTransformerEncoder(50, 5, 10, 100, 2, 'isab', 3)
st_dec = nn.SetTransformerDecoder(50, 5, 10, 100, 2, 4)
print(st_enc_0, st_enc_1, st_dec)
# test#1: basic
h0 = th.rand(g.number_of_nodes(), 50)
h1 = st_enc_0(h0, g)
assert h1.shape == h0.shape
h1 = st_enc_1(h0, g)
assert h1.shape == h0.shape
h2 = st_dec(h1, g)
assert h2.shape[0] == 200 and h2.dim() == 1
# test#2: batched graph
g1 = dgl.DGLGraph(nx.path_graph(5))
g2 = dgl.DGLGraph(nx.path_graph(10))
bg = dgl.batch([g, g1, g2])
h0 = th.rand(bg.number_of_nodes(), 50)
h1 = st_enc_0(h0, bg)
assert h1.shape == h0.shape
h1 = st_enc_1(h0, bg)
assert h1.shape == h0.shape
h2 = st_dec(h1, bg)
assert h2.shape[0] == 3 and h2.shape[1] == 200 and h2.dim() == 2
