|
1 | 1 | @testset "HeteroGraphConv" begin
|
2 | 2 | d, n = 3, 5
|
3 | 3 | g = rand_bipartite_heterograph(n, 2*n, 15)
|
| 4 | + hg = rand_bipartite_heterograph((2,3), 6) |
4 | 5 |
|
5 | 6 | model = HeteroGraphConv([(:A,:to,:B) => GraphConv(d => d),
|
6 | 7 | (:B,:to,:A) => GraphConv(d => d)])
|
|
93 | 94 | end
|
94 | 95 |
|
95 | 96 | @testset "CGConv" begin
|
96 |
| - g = rand_bipartite_heterograph((2,3), 6) |
97 | 97 | x = (A = rand(Float32, 4,2), B = rand(Float32, 4, 3))
|
98 | 98 | layers = HeteroGraphConv( (:A, :to, :B) => CGConv(4 => 2, relu),
|
99 | 99 | (:B, :to, :A) => CGConv(4 => 2, relu));
|
100 |
| - y = layers(g, x); |
| 100 | + y = layers(hg, x); |
101 | 101 | @test size(y.A) == (2,2) && size(y.B) == (2,3)
|
102 | 102 | end
|
103 | 103 |
|
104 | 104 | @testset "EdgeConv" begin
|
105 |
| - g = rand_bipartite_heterograph((2,3), 6) |
106 | 105 | x = (A = rand(Float32, 4,2), B = rand(Float32, 4, 3))
|
107 | 106 | layers = HeteroGraphConv( (:A, :to, :B) => EdgeConv(Dense(2 * 4, 2), aggr = +),
|
108 | 107 | (:B, :to, :A) => EdgeConv(Dense(2 * 4, 2), aggr = +));
|
109 |
| - y = layers(g, x); |
| 108 | + y = layers(hg, x); |
110 | 109 | @test size(y.A) == (2,2) && size(y.B) == (2,3)
|
111 | 110 | end
|
112 | 111 |
|
|
0 commit comments