support batch

Author: yuehhua

src/layers/group_conv.jl

@@ -23,7 +23,7 @@ EEquivGraphConv(ϕ_edge=Dense(10 => 5), ϕ_x=Dense(5 => 2), ϕ_h=Dense(8 => 5))
 
 See also [`WithGraph`](@ref) for training layer with static graph and [`EEquivGraphPE`](@ref) for positional encoding.
 """
-struct EEquivGraphConv{X,E,H}
+struct EEquivGraphConv{X,E,H} <: AbstractGraphLayer
     pe::X
     nn_edge::E
     nn_h::H
@@ -61,6 +61,15 @@ function(egnn::EEquivGraphConv)(fg::AbstractFeaturedGraph)
     return ConcreteFeaturedGraph(fg, nf=V, pf=X)
 end
 
+# For static graph
+function(l::EEquivGraphConv)(el::NamedTuple, H::AbstractArray, E::AbstractArray, X::AbstractArray)
+    GraphSignals.check_num_nodes(el.N, H)
+    GraphSignals.check_num_nodes(el.N, X)
+    GraphSignals.check_num_edges(el.E, E)
+    _, V, X = propagate(l, el, E, H, X, +)
+    return V, X
+end
+
 function Base.show(io::IO, l::EEquivGraphConv)
     print(io, "EEquivGraphConv(ϕ_edge=", l.nn_edge)
     print(io, ", ϕ_x=", l.pe.nn)
@@ -96,3 +105,6 @@ function propagate(l::EEquivGraphConv, el::NamedTuple, E, V, X, aggr)
     V = update(l, Ē, V)
     return E, V, X
 end
+
+WithGraph(fg::AbstractFeaturedGraph, l::EEquivGraphConv) = WithGraph(to_namedtuple(fg), l)
+(wg::WithGraph{<:EEquivGraphConv})(args...) = wg.layer(wg.graph, args...)

test/layers/group_conv.jl

@@ -32,5 +32,18 @@
             g = Zygote.gradient(() -> sum(node_feature(egnn(fg))), Flux.params(egnn))
             @test length(g.grads) == 8
         end
+
+        @testset "layer with static graph" begin
+            nf = rand(T, in_channel, N, batch_size)
+            ef = rand(T, in_channel_edge, E, batch_size)
+            pf = rand(T, pos_dim, N, batch_size)
+            l = WithGraph(fg, EEquivGraphConv(in_channel=>out_channel, pos_dim, in_channel_edge))
+            H, Y = l(nf, ef, pf)
+            @test size(H) == (out_channel, N, batch_size)
+            @test size(Y) == (pos_dim, N, batch_size)
+
+            g = Zygote.gradient(() -> sum(l(nf, ef, pf)[1]), Flux.params(l))
+            @test length(g.grads) == 6
+        end
     end
 end