Only uploading Gate activation and uploading Lux TGCN custom non linear functions

Author: juanmigutierrez

GNNLux/src/layers/temporalconv.jl

@@ -33,10 +33,21 @@ LuxCore.apply(m::GNNContainerLayer, g, x, ps, st) = m(g, x, ps, st)
     init_state::Function
 end
 
-function TGCNCell(ch::Pair{Int, Int}; use_bias = true, init_weight = glorot_uniform, init_state = zeros32, init_bias = zeros32, add_self_loops = false, use_edge_weight = true) 
+function TGCNCell(ch::Pair{Int, Int}; 
+                 use_bias = true, 
+                 init_weight = glorot_uniform, 
+                 init_state = zeros32, 
+                 init_bias = zeros32, 
+                 add_self_loops = false, 
+                 use_edge_weight = true,
+                 gate_activation = sigmoid) 
     in_dims, out_dims = ch
-    conv = GCNConv(ch, sigmoid; init_weight, init_bias, use_bias, add_self_loops, use_edge_weight)
-    gru = Lux.GRUCell(out_dims => out_dims; use_bias, init_weight = (init_weight, init_weight, init_weight), init_bias = (init_bias, init_bias, init_bias), init_state = init_state)
+    conv = GCNConv(ch, gate_activation; init_weight, init_bias, use_bias, add_self_loops, use_edge_weight)
+    gru = Lux.GRUCell(out_dims => out_dims; 
+                      use_bias, 
+                      init_weight = (init_weight, init_weight, init_weight), 
+                      init_bias = (init_bias, init_bias, init_bias), 
+                      init_state = init_state)
     return TGCNCell(in_dims, out_dims, conv, gru, init_state)
 end
 

GNNLux/test/layers/temporalconv.jl

@@ -56,4 +56,12 @@
         loss = (tx, ps) -> sum(sum(first(l(tg, tx, ps, st))))
         test_gradients(loss, tx, ps; atol=1.0f-2, rtol=1.0f-2, skip_backends=[AutoReverseDiff(), AutoTracker(), AutoForwardDiff(), AutoEnzyme()])
     end
+
+    @testset "TGCN with Custom Activations" begin
+        l = TGCN(3=>3, gate_activation = Lux.relu)
+        ps = LuxCore.initialparameters(rng, l)
+        st = LuxCore.initialstates(rng, l)
+        loss = (x, ps) -> sum(first(l(g, x, ps, st)))
+        test_gradients(loss, x, ps; atol=1.0f-2, rtol=1.0f-2, skip_backends=[AutoReverseDiff(), AutoTracker(), AutoForwardDiff(), AutoEnzyme()])
+    end
 end

GraphNeuralNetworks/src/layers/temporalconv.jl

@@ -826,14 +826,13 @@ Flux.@layer :noexpand TGCNCell
 
 function TGCNCell((in, out)::Pair{Int, Int}; 
                   gate_activation = sigmoid,
-                  hidden_activation = tanh,
                   kws...)
     conv_z = GNNChain(GCNConv(in => out, relu; kws...), GCNConv(out => out; kws...))
     dense_z = Dense(2*out => out, gate_activation)
     conv_r = GNNChain(GCNConv(in => out, relu; kws...), GCNConv(out => out; kws...))
     dense_r = Dense(2*out => out, gate_activation)
     conv_h = GNNChain(GCNConv(in => out, relu; kws...), GCNConv(out => out; kws...))
-    dense_h = Dense(2*out => out, hidden_activation)
+    dense_h = Dense(2*out => out, tanh)
     return TGCNCell(in, out, conv_z, dense_z, conv_r, dense_r, conv_h, dense_h)
 end
 
@@ -897,6 +896,6 @@ julia> size(y) # (d_out, timesteps, num_nodes)
 (3, 5, 5)
 ```
 """
-TGCN(args...; gate_activation = sigmoid, hidden_activation = tanh, kws...) = 
-    GNNRecurrence(TGCNCell(args...; gate_activation, hidden_activation, kws...))
+TGCN(args...; gate_activation = sigmoid, kws...) = 
+    GNNRecurrence(TGCNCell(args...; gate_activation, kws...))
 

GraphNeuralNetworks/test/layers/temporalconv.jl

@@ -59,7 +59,7 @@ end
     using .TemporalConvTestModule, .TestModule
     using Flux: relu, sigmoid
     layer_default = TGCN(in_channel => out_channel)
-    layer_custom = TGCN(in_channel => out_channel, gate_activation = relu, hidden_activation = relu)
+    layer_custom = TGCN(in_channel => out_channel, gate_activation = relu)
     x = rand(Float32, in_channel, timesteps, g.num_nodes)
     y_default = layer_default(g, x)
     y_custom = layer_custom(g, x)