From ee24b650af455bfd086eaaacc799346ed72f52a1 Mon Sep 17 00:00:00 2001
From: juanmigutierrez <recuperacionmiguel@outlook.es>
Date: Fri, 28 Mar 2025 23:04:38 +0100
Subject: [PATCH 1/9] Adding TGCN with parameters for custom activations
 functions in hidden and gate activations

---
 .../src/layers/temporalconv.jl                | 23 +++++++++++++------
 .../test/layers/temporalconv.jl               | 22 ++++++++++++++++++
 2 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/GraphNeuralNetworks/src/layers/temporalconv.jl b/GraphNeuralNetworks/src/layers/temporalconv.jl
index 2872529d6..4d5b16dc0 100644
--- a/GraphNeuralNetworks/src/layers/temporalconv.jl
+++ b/GraphNeuralNetworks/src/layers/temporalconv.jl
@@ -824,13 +824,16 @@ end
 
 Flux.@layer :noexpand TGCNCell
 
-function TGCNCell((in, out)::Pair{Int, Int}; kws...)
+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, sigmoid)
+    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, sigmoid)
+    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, tanh)
+    dense_h = Dense(2*out => out, hidden_activation)
     return TGCNCell(in, out, conv_z, dense_z, conv_r, dense_r, conv_h, dense_h)
 end
 
@@ -858,13 +861,18 @@ function Base.show(io::IO, cell::TGCNCell)
 end
 
 """
-    TGCN(args...; kws...)
+    TGCN(args...; gate_activation = sigmoid, hidden_activation = tanh, kws...)
 
 Construct a recurrent layer corresponding to the [`TGCNCell`](@ref) cell.
 
 The arguments are passed to the [`TGCNCell`](@ref) constructor.
 See [`GNNRecurrence`](@ref) for more details.
 
+# Additional Parameters
+
+- `gate_activation`: Activation function for the gate mechanisms. Default `sigmoid`.
+- `hidden_activation`: Activation function for the hidden state update. Default `tanh`.
+
 # Examples
 
 ```jldoctest
@@ -878,7 +886,7 @@ julia> g = rand_graph(num_nodes, num_edges);
 
 julia> x = rand(Float32, d_in, timesteps, num_nodes);
 
-julia> layer = TGCN(d_in => d_out)
+julia> layer = TGCN(d_in => d_out, hidden_activation = relu)
 GNNRecurrence(
   TGCNCell(2 => 3),                     # 126 parameters
 )                   # Total: 18 arrays, 126 parameters, 1.469 KiB.
@@ -889,5 +897,6 @@ julia> size(y) # (d_out, timesteps, num_nodes)
 (3, 5, 5)
 ```
 """
-TGCN(args...; kws...) = GNNRecurrence(TGCNCell(args...; kws...))
+TGCN(args...; gate_activation = sigmoid, hidden_activation = tanh, kws...) = 
+    GNNRecurrence(TGCNCell(args...; gate_activation, hidden_activation, kws...))
 
diff --git a/GraphNeuralNetworks/test/layers/temporalconv.jl b/GraphNeuralNetworks/test/layers/temporalconv.jl
index f8d96c0e2..041e36147 100644
--- a/GraphNeuralNetworks/test/layers/temporalconv.jl
+++ b/GraphNeuralNetworks/test/layers/temporalconv.jl
@@ -55,6 +55,28 @@ end
     test_gradients(model, g, x, rtol = RTOL_HIGH, atol = ATOL_LOW)
 end
 
+@testitem "TGCN with custom activations" setup=[TemporalConvTestModule, TestModule] begin
+    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)
+    x = rand(Float32, in_channel, timesteps, g.num_nodes)
+    y_default = layer_default(g, x)
+    y_custom = layer_custom(g, x)
+    @test layer_default isa GNNRecurrence
+    @test layer_custom isa GNNRecurrence
+    @test size(y_default) == size(y_custom)
+    # Outputs should be different due to different activations
+    @test y_default != y_custom
+    test_gradients(layer_custom, g, x, rtol = RTOL_HIGH)
+
+    # interplay with GNNChain
+    model = GNNChain(TGCN(in_channel => out_channel), Dense(out_channel, 1))
+    y = model(g, x)
+    @test size(y) == (1, timesteps, g.num_nodes)
+    test_gradients(model, g, x, rtol = RTOL_HIGH, atol = ATOL_LOW)
+end
+
 @testitem "GConvLSTMCell" setup=[TemporalConvTestModule, TestModule] begin
     using .TemporalConvTestModule, .TestModule
     cell = GConvLSTMCell(in_channel => out_channel, 2)

From be10c3484b8e761eaa5c0a6ff9b6f3eb62b19579 Mon Sep 17 00:00:00 2001
From: juanmigutierrez <recuperacionmiguel@outlook.es>
Date: Fri, 28 Mar 2025 23:47:50 +0100
Subject: [PATCH 2/9] Only uploading Gate activation and uploading Lux TGCN
 custom non linear functions

---
 GNNLux/src/layers/temporalconv.jl               | 17 ++++++++++++++---
 GNNLux/test/layers/temporalconv.jl              |  8 ++++++++
 GraphNeuralNetworks/src/layers/temporalconv.jl  |  7 +++----
 GraphNeuralNetworks/test/layers/temporalconv.jl |  2 +-
 4 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/GNNLux/src/layers/temporalconv.jl b/GNNLux/src/layers/temporalconv.jl
index e58901774..73c90457f 100644
--- a/GNNLux/src/layers/temporalconv.jl
+++ b/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
 
diff --git a/GNNLux/test/layers/temporalconv.jl b/GNNLux/test/layers/temporalconv.jl
index 8de93efa7..4a231a939 100644
--- a/GNNLux/test/layers/temporalconv.jl
+++ b/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
\ No newline at end of file
diff --git a/GraphNeuralNetworks/src/layers/temporalconv.jl b/GraphNeuralNetworks/src/layers/temporalconv.jl
index 4d5b16dc0..e6dea936e 100644
--- a/GraphNeuralNetworks/src/layers/temporalconv.jl
+++ b/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...))
 
diff --git a/GraphNeuralNetworks/test/layers/temporalconv.jl b/GraphNeuralNetworks/test/layers/temporalconv.jl
index 041e36147..6d9b2da28 100644
--- a/GraphNeuralNetworks/test/layers/temporalconv.jl
+++ b/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)

From fc2941b16c88622f48109251b55cd7e86565ff48 Mon Sep 17 00:00:00 2001
From: juanmigutierrez <recuperacionmiguel@outlook.es>
Date: Sat, 29 Mar 2025 00:12:37 +0100
Subject: [PATCH 3/9] Adding gate_activation=relu for test GNNChain in TGCN and
 fixing documentation

---
 GraphNeuralNetworks/src/layers/temporalconv.jl  | 5 ++---
 GraphNeuralNetworks/test/layers/temporalconv.jl | 2 +-
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/GraphNeuralNetworks/src/layers/temporalconv.jl b/GraphNeuralNetworks/src/layers/temporalconv.jl
index e6dea936e..222be4ecd 100644
--- a/GraphNeuralNetworks/src/layers/temporalconv.jl
+++ b/GraphNeuralNetworks/src/layers/temporalconv.jl
@@ -860,7 +860,7 @@ function Base.show(io::IO, cell::TGCNCell)
 end
 
 """
-    TGCN(args...; gate_activation = sigmoid, hidden_activation = tanh, kws...)
+    TGCN(args...; gate_activation = sigmoid, kws...)
 
 Construct a recurrent layer corresponding to the [`TGCNCell`](@ref) cell.
 
@@ -870,7 +870,6 @@ See [`GNNRecurrence`](@ref) for more details.
 # Additional Parameters
 
 - `gate_activation`: Activation function for the gate mechanisms. Default `sigmoid`.
-- `hidden_activation`: Activation function for the hidden state update. Default `tanh`.
 
 # Examples
 
@@ -885,7 +884,7 @@ julia> g = rand_graph(num_nodes, num_edges);
 
 julia> x = rand(Float32, d_in, timesteps, num_nodes);
 
-julia> layer = TGCN(d_in => d_out, hidden_activation = relu)
+julia> layer = TGCN(d_in => d_out, gate_activation = relu)
 GNNRecurrence(
   TGCNCell(2 => 3),                     # 126 parameters
 )                   # Total: 18 arrays, 126 parameters, 1.469 KiB.
diff --git a/GraphNeuralNetworks/test/layers/temporalconv.jl b/GraphNeuralNetworks/test/layers/temporalconv.jl
index 6d9b2da28..0b94476c2 100644
--- a/GraphNeuralNetworks/test/layers/temporalconv.jl
+++ b/GraphNeuralNetworks/test/layers/temporalconv.jl
@@ -71,7 +71,7 @@ end
     test_gradients(layer_custom, g, x, rtol = RTOL_HIGH)
 
     # interplay with GNNChain
-    model = GNNChain(TGCN(in_channel => out_channel), Dense(out_channel, 1))
+    model = GNNChain(TGCN(in_channel => out_channel, gate_activation=relu), Dense(out_channel, 1))
     y = model(g, x)
     @test size(y) == (1, timesteps, g.num_nodes)
     test_gradients(model, g, x, rtol = RTOL_HIGH, atol = ATOL_LOW)

From ebee230b45dbcd7a0cf0107e1049ae40d8544d4a Mon Sep 17 00:00:00 2001
From: juanmigutierrez <recuperacionmiguel@outlook.es>
Date: Sat, 29 Mar 2025 01:03:02 +0100
Subject: [PATCH 4/9] try to fix doctest

---
 GraphNeuralNetworks/src/layers/temporalconv.jl | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/GraphNeuralNetworks/src/layers/temporalconv.jl b/GraphNeuralNetworks/src/layers/temporalconv.jl
index 222be4ecd..422563e41 100644
--- a/GraphNeuralNetworks/src/layers/temporalconv.jl
+++ b/GraphNeuralNetworks/src/layers/temporalconv.jl
@@ -874,6 +874,8 @@ See [`GNNRecurrence`](@ref) for more details.
 # Examples
 
 ```jldoctest
+julia> using Flux  # Ensure relu is available
+
 julia> num_nodes, num_edges = 5, 10;
 
 julia> d_in, d_out = 2, 3;

From b204a0a16f25c8803329c1bbb6486bd2045b1041 Mon Sep 17 00:00:00 2001
From: juanmigutierrez <recuperacionmiguel@outlook.es>
Date: Sun, 30 Mar 2025 18:11:27 +0200
Subject: [PATCH 5/9] Chaging f(A,x) activation non linear function of TGCN and
 the corresponding tests

---
 GNNLux/src/layers/temporalconv.jl             | 15 ++++++----
 GNNLux/test/layers/temporalconv.jl            | 14 +++++++++
 .../src/layers/temporalconv.jl                | 23 +++++++++-----
 .../test/layers/temporalconv.jl               | 30 +++++++++++++++++++
 4 files changed, 69 insertions(+), 13 deletions(-)

diff --git a/GNNLux/src/layers/temporalconv.jl b/GNNLux/src/layers/temporalconv.jl
index e58901774..d9e5d4c55 100644
--- a/GNNLux/src/layers/temporalconv.jl
+++ b/GNNLux/src/layers/temporalconv.jl
@@ -33,9 +33,11 @@ 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, act = relu) 
     in_dims, out_dims = ch
-    conv = GCNConv(ch, sigmoid; init_weight, init_bias, use_bias, add_self_loops, use_edge_weight)
+    conv = GCNConv(ch, act; 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
@@ -57,7 +59,7 @@ function Base.show(io::IO, tgcn::TGCNCell)
 end
 
 """
-    TGCN(in => out; use_bias = true, init_weight = glorot_uniform, init_state = zeros32, init_bias = zeros32, add_self_loops = false, use_edge_weight = true) 
+    TGCN(in => out; use_bias = true, init_weight = glorot_uniform, init_state = zeros32, init_bias = zeros32, add_self_loops = false, use_edge_weight = true, act = sigmoid) 
 
 Temporal Graph Convolutional Network (T-GCN) recurrent layer from the paper [T-GCN: A Temporal Graph Convolutional Network for Traffic Prediction](https://arxiv.org/pdf/1811.05320.pdf).
 
@@ -76,7 +78,7 @@ Performs a layer of GCNConv to model spatial dependencies, followed by a Gated R
                      If `add_self_loops=true` the new weights will be set to 1. 
                      This option is ignored if the `edge_weight` is explicitly provided in the forward pass.
                      Default `false`.
-
+- `act`: Activation function used in the GCNConv layer. Default `relu`.
 
 
 # Examples
@@ -91,9 +93,12 @@ rng = Random.default_rng()
 g = rand_graph(rng, 5, 10)
 x = rand(rng, Float32, 2, 5)
 
-# create TGCN layer
+# create TGCN layer with default activation (relu)
 tgcn = TGCN(2 => 6)
 
+# create TGCN layer with custom activation
+tgcn_relu = TGCN(2 => 6, act = sigmoid)
+
 # setup layer
 ps, st = LuxCore.setup(rng, tgcn)
 
diff --git a/GNNLux/test/layers/temporalconv.jl b/GNNLux/test/layers/temporalconv.jl
index 8de93efa7..1fa3354e4 100644
--- a/GNNLux/test/layers/temporalconv.jl
+++ b/GNNLux/test/layers/temporalconv.jl
@@ -10,11 +10,25 @@
     tx = [x for _ in 1:5]
 
     @testset "TGCN" begin
+        # Test with default activation (relu)
         l = TGCN(3=>3)
         ps = LuxCore.initialparameters(rng, l)
         st = LuxCore.initialstates(rng, l)
+        y1, _ = l(g, x, ps, st)
         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()])
+        
+        # Test with custom activation (sigmoid)
+        l_relu = TGCN(3=>3, act = sigmoid)
+        ps_relu = LuxCore.initialparameters(rng, l_relu)
+        st_relu = LuxCore.initialstates(rng, l_relu)
+        y2, _ = l_relu(g, x, ps_relu, st_relu)
+        
+        # Outputs should be different with different activation functions
+        @test !isapprox(y1, y2, rtol=1.0f-2)
+        
+        loss_relu = (x, ps) -> sum(first(l_relu(g, x, ps, st_relu)))
+        test_gradients(loss_relu, x, ps_relu; atol=1.0f-2, rtol=1.0f-2, skip_backends=[AutoReverseDiff(), AutoTracker(), AutoForwardDiff(), AutoEnzyme()])
     end
 
     @testset "A3TGCN" begin
diff --git a/GraphNeuralNetworks/src/layers/temporalconv.jl b/GraphNeuralNetworks/src/layers/temporalconv.jl
index 2872529d6..827d313d7 100644
--- a/GraphNeuralNetworks/src/layers/temporalconv.jl
+++ b/GraphNeuralNetworks/src/layers/temporalconv.jl
@@ -758,7 +758,7 @@ EvolveGCNO(args...; kws...) = GNNRecurrence(EvolveGCNOCell(args...; kws...))
 
 
 """
-    TGCNCell(in => out; kws...)
+    TGCNCell(in => out; act = relu, kws...)
 
 Recurrent graph convolutional cell from the paper
 [T-GCN: A Temporal Graph Convolutional
@@ -824,12 +824,14 @@ end
 
 Flux.@layer :noexpand TGCNCell
 
-function TGCNCell((in, out)::Pair{Int, Int}; kws...)
-    conv_z = GNNChain(GCNConv(in => out, relu; kws...), GCNConv(out => out; kws...))
+function TGCNCell((in, out)::Pair{Int, Int}; 
+    act = relu,
+    kws...)
+    conv_z = GNNChain(GCNConv(in => out, act; kws...), GCNConv(out => out; kws...))
     dense_z = Dense(2*out => out, sigmoid)
-    conv_r = GNNChain(GCNConv(in => out, relu; kws...), GCNConv(out => out; kws...))
+    conv_r = GNNChain(GCNConv(in => out, act; kws...), GCNConv(out => out; kws...))
     dense_r = Dense(2*out => out, sigmoid)
-    conv_h = GNNChain(GCNConv(in => out, relu; kws...), GCNConv(out => out; kws...))
+    conv_h = GNNChain(GCNConv(in => out, act; kws...), GCNConv(out => out; kws...))
     dense_h = Dense(2*out => out, tanh)
     return TGCNCell(in, out, conv_z, dense_z, conv_r, dense_r, conv_h, dense_h)
 end
@@ -858,7 +860,7 @@ function Base.show(io::IO, cell::TGCNCell)
 end
 
 """
-    TGCN(args...; kws...)
+    TGCN(args...; act = relu, kws...)
 
 Construct a recurrent layer corresponding to the [`TGCNCell`](@ref) cell.
 
@@ -876,9 +878,14 @@ julia> timesteps = 5;
 
 julia> g = rand_graph(num_nodes, num_edges);
 
-julia> x = rand(Float32, d_in, timesteps, num_nodes);
+julia> x = rand(Float32, d_in, timesteps, g.num_nodes);
+
+julia> layer = TGCN(d_in => d_out)  # Default activation (relu)
+GNNRecurrence(
+  TGCNCell(2 => 3),                     # 126 parameters
+)                   # Total: 18 arrays, 126 parameters, 1.469 KiB.
 
-julia> layer = TGCN(d_in => d_out)
+julia> layer_tanh = TGCN(d_in => d_out, act = tanh)  # Custom activation
 GNNRecurrence(
   TGCNCell(2 => 3),                     # 126 parameters
 )                   # Total: 18 arrays, 126 parameters, 1.469 KiB.
diff --git a/GraphNeuralNetworks/test/layers/temporalconv.jl b/GraphNeuralNetworks/test/layers/temporalconv.jl
index f8d96c0e2..93d6b0082 100644
--- a/GraphNeuralNetworks/test/layers/temporalconv.jl
+++ b/GraphNeuralNetworks/test/layers/temporalconv.jl
@@ -25,6 +25,8 @@ end
 
 @testitem "TGCNCell" setup=[TemporalConvTestModule, TestModule] begin
     using .TemporalConvTestModule, .TestModule
+    
+    # Test with default activation function
     cell = GraphNeuralNetworks.TGCNCell(in_channel => out_channel)
     y, h = cell(g, g.x)
     @test y === h
@@ -33,10 +35,25 @@ end
     test_gradients(cell, g, g.x, loss=cell_loss, rtol=RTOL_HIGH)
     # with initial state
     test_gradients(cell, g, g.x, h, loss=cell_loss, rtol=RTOL_HIGH)
+    
+    # Test with custom activation function
+    custom_activation = tanh
+    cell_custom = GraphNeuralNetworks.TGCNCell(in_channel => out_channel, act = custom_activation)
+    y_custom, h_custom = cell_custom(g, g.x)
+    @test y_custom === h_custom
+    @test size(h_custom) == (out_channel, g.num_nodes)
+    # Test that outputs differ when using different activation functions
+    @test !isapprox(y, y_custom, rtol=RTOL_HIGH)
+    # with no initial state
+    test_gradients(cell_custom, g, g.x, loss=cell_loss, rtol=RTOL_HIGH)
+    # with initial state
+    test_gradients(cell_custom, g, g.x, h_custom, loss=cell_loss, rtol=RTOL_HIGH)
 end
 
 @testitem "TGCN" setup=[TemporalConvTestModule, TestModule] begin
     using .TemporalConvTestModule, .TestModule
+    
+    # Test with default activation function
     layer = TGCN(in_channel => out_channel)
     x = rand(Float32, in_channel, timesteps, g.num_nodes)
     state0 = rand(Float32, out_channel, g.num_nodes)
@@ -48,6 +65,19 @@ end
     # with initial state
     test_gradients(layer, g, x, state0, rtol = RTOL_HIGH)
 
+    # Test with custom activation function
+    custom_activation = tanh
+    layer_custom = TGCN(in_channel => out_channel, act = custom_activation)
+    y_custom = layer_custom(g, x)
+    @test layer_custom isa GNNRecurrence
+    @test size(y_custom) == (out_channel, timesteps, g.num_nodes)
+    # Test that outputs differ when using different activation functions
+    @test !isapprox(y, y_custom, rtol = RTOL_HIGH)
+    # with no initial state
+    test_gradients(layer_custom, g, x, rtol = RTOL_HIGH)
+    # with initial state
+    test_gradients(layer_custom, g, x, state0, rtol = RTOL_HIGH)
+
     # interplay with GNNChain
     model = GNNChain(TGCN(in_channel => out_channel), Dense(out_channel, 1))
     y = model(g, x)

From 883f04d6ea417e152286f18eb4a7e4b9100f7d0e Mon Sep 17 00:00:00 2001
From: juanmigutierrez <recuperacionmiguel@outlook.es>
Date: Sun, 30 Mar 2025 19:09:10 +0200
Subject: [PATCH 6/9] leave sigmoid as before

---
 GNNLux/src/layers/temporalconv.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/GNNLux/src/layers/temporalconv.jl b/GNNLux/src/layers/temporalconv.jl
index d9e5d4c55..7d25c218a 100644
--- a/GNNLux/src/layers/temporalconv.jl
+++ b/GNNLux/src/layers/temporalconv.jl
@@ -35,7 +35,7 @@ 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, act = relu) 
+                    use_edge_weight = true, act = sigmoid) 
     in_dims, out_dims = ch
     conv = GCNConv(ch, act; 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)
@@ -59,7 +59,7 @@ function Base.show(io::IO, tgcn::TGCNCell)
 end
 
 """
-    TGCN(in => out; use_bias = true, init_weight = glorot_uniform, init_state = zeros32, init_bias = zeros32, add_self_loops = false, use_edge_weight = true, act = sigmoid) 
+    TGCN(in => out; use_bias = true, init_weight = glorot_uniform, init_state = zeros32, init_bias = zeros32, add_self_loops = false, use_edge_weight = true, act = relu) 
 
 Temporal Graph Convolutional Network (T-GCN) recurrent layer from the paper [T-GCN: A Temporal Graph Convolutional Network for Traffic Prediction](https://arxiv.org/pdf/1811.05320.pdf).
 

From 560036842b9e614adc2ef8660180b245412a44eb Mon Sep 17 00:00:00 2001
From: juanmigutierrez <recuperacionmiguel@outlook.es>
Date: Sun, 30 Mar 2025 19:14:41 +0200
Subject: [PATCH 7/9] leave docs of GNNLux as before

---
 GNNLux/src/layers/temporalconv.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/GNNLux/src/layers/temporalconv.jl b/GNNLux/src/layers/temporalconv.jl
index 7d25c218a..df9ba4bc4 100644
--- a/GNNLux/src/layers/temporalconv.jl
+++ b/GNNLux/src/layers/temporalconv.jl
@@ -59,7 +59,7 @@ function Base.show(io::IO, tgcn::TGCNCell)
 end
 
 """
-    TGCN(in => out; use_bias = true, init_weight = glorot_uniform, init_state = zeros32, init_bias = zeros32, add_self_loops = false, use_edge_weight = true, act = relu) 
+    TGCN(in => out; use_bias = true, init_weight = glorot_uniform, init_state = zeros32, init_bias = zeros32, add_self_loops = false, use_edge_weight = true, act = sigmoid) 
 
 Temporal Graph Convolutional Network (T-GCN) recurrent layer from the paper [T-GCN: A Temporal Graph Convolutional Network for Traffic Prediction](https://arxiv.org/pdf/1811.05320.pdf).
 
@@ -78,7 +78,7 @@ Performs a layer of GCNConv to model spatial dependencies, followed by a Gated R
                      If `add_self_loops=true` the new weights will be set to 1. 
                      This option is ignored if the `edge_weight` is explicitly provided in the forward pass.
                      Default `false`.
-- `act`: Activation function used in the GCNConv layer. Default `relu`.
+- `act`: Activation function used in the GCNConv layer. Default `sigmoid`.
 
 
 # Examples
@@ -93,11 +93,11 @@ rng = Random.default_rng()
 g = rand_graph(rng, 5, 10)
 x = rand(rng, Float32, 2, 5)
 
-# create TGCN layer with default activation (relu)
+# create TGCN layer with default activation (sigmoid)
 tgcn = TGCN(2 => 6)
 
 # create TGCN layer with custom activation
-tgcn_relu = TGCN(2 => 6, act = sigmoid)
+tgcn_relu = TGCN(2 => 6, act = relu)
 
 # setup layer
 ps, st = LuxCore.setup(rng, tgcn)

From 7c8dfeb322a8618180bc9f280bf89bc0d1207902 Mon Sep 17 00:00:00 2001
From: juanmigutierrez <recuperacionmiguel@outlook.es>
Date: Sun, 30 Mar 2025 20:17:19 +0200
Subject: [PATCH 8/9] fix documentation and test

---
 GNNLux/src/layers/temporalconv.jl              | 2 +-
 GNNLux/test/layers/temporalconv.jl             | 6 +++---
 GraphNeuralNetworks/src/layers/temporalconv.jl | 4 ++--
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/GNNLux/src/layers/temporalconv.jl b/GNNLux/src/layers/temporalconv.jl
index df9ba4bc4..d6036cd8f 100644
--- a/GNNLux/src/layers/temporalconv.jl
+++ b/GNNLux/src/layers/temporalconv.jl
@@ -93,7 +93,7 @@ rng = Random.default_rng()
 g = rand_graph(rng, 5, 10)
 x = rand(rng, Float32, 2, 5)
 
-# create TGCN layer with default activation (sigmoid)
+# create TGCN layer 
 tgcn = TGCN(2 => 6)
 
 # create TGCN layer with custom activation
diff --git a/GNNLux/test/layers/temporalconv.jl b/GNNLux/test/layers/temporalconv.jl
index 1fa3354e4..794775f88 100644
--- a/GNNLux/test/layers/temporalconv.jl
+++ b/GNNLux/test/layers/temporalconv.jl
@@ -10,7 +10,7 @@
     tx = [x for _ in 1:5]
 
     @testset "TGCN" begin
-        # Test with default activation (relu)
+        # Test with default activation (sigmoid)
         l = TGCN(3=>3)
         ps = LuxCore.initialparameters(rng, l)
         st = LuxCore.initialstates(rng, l)
@@ -18,8 +18,8 @@
         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()])
         
-        # Test with custom activation (sigmoid)
-        l_relu = TGCN(3=>3, act = sigmoid)
+        # Test with custom activation (relu)
+        l_relu = TGCN(3=>3, act = relu)
         ps_relu = LuxCore.initialparameters(rng, l_relu)
         st_relu = LuxCore.initialstates(rng, l_relu)
         y2, _ = l_relu(g, x, ps_relu, st_relu)
diff --git a/GraphNeuralNetworks/src/layers/temporalconv.jl b/GraphNeuralNetworks/src/layers/temporalconv.jl
index ce040dfac..18150a97c 100644
--- a/GraphNeuralNetworks/src/layers/temporalconv.jl
+++ b/GraphNeuralNetworks/src/layers/temporalconv.jl
@@ -869,12 +869,12 @@ See [`GNNRecurrence`](@ref) for more details.
 
 # Additional Parameters
 
-- `gate_activation`: Activation function for the gate mechanisms. Default `sigmoid`.
+- `act`: Activation function for the GCNConv layers. Default `relu`.
 
 # Examples
 
 ```jldoctest
-julia> using Flux  # Ensure relu is available
+julia> using Flux  # Ensure activation functions are available
 
 julia> num_nodes, num_edges = 5, 10;
 

From 75fed448e0addb80318dcfc58af833daae7611f0 Mon Sep 17 00:00:00 2001
From: juanmigutierrez <recuperacionmiguel@outlook.es>
Date: Wed, 9 Apr 2025 18:54:33 +0200
Subject: [PATCH 9/9] Solving recomendations

---
 GraphNeuralNetworks/src/layers/temporalconv.jl | 12 ++++--------
 1 file changed, 4 insertions(+), 8 deletions(-)

diff --git a/GraphNeuralNetworks/src/layers/temporalconv.jl b/GraphNeuralNetworks/src/layers/temporalconv.jl
index 18150a97c..5516d8015 100644
--- a/GraphNeuralNetworks/src/layers/temporalconv.jl
+++ b/GraphNeuralNetworks/src/layers/temporalconv.jl
@@ -758,7 +758,7 @@ EvolveGCNO(args...; kws...) = GNNRecurrence(EvolveGCNOCell(args...; kws...))
 
 
 """
-    TGCNCell(in => out; act = relu, kws...)
+    TGCNCell(in => out, act = relu, kws...)
 
 Recurrent graph convolutional cell from the paper
 [T-GCN: A Temporal Graph Convolutional
@@ -860,17 +860,13 @@ function Base.show(io::IO, cell::TGCNCell)
 end
 
 """
-    TGCN(args...; act = relu, kws...)
+    TGCN(args...; kws...)
 
 Construct a recurrent layer corresponding to the [`TGCNCell`](@ref) cell.
 
 The arguments are passed to the [`TGCNCell`](@ref) constructor.
 See [`GNNRecurrence`](@ref) for more details.
 
-# Additional Parameters
-
-- `act`: Activation function for the GCNConv layers. Default `relu`.
-
 # Examples
 
 ```jldoctest
@@ -902,6 +898,6 @@ julia> size(y) # (d_out, timesteps, num_nodes)
 (3, 5, 5)
 ```
 """
-TGCN(args...; act = relu, kws...) = 
-    GNNRecurrence(TGCNCell(args...; act, kws...))
+TGCN(args...; kws...) = 
+    GNNRecurrence(TGCNCell(args...; kws...))