Merge branch 'main' into advance

2024-05-24 19:42:03 +08:00
parent 40b7d8360c 1fe5d5fca2
commit 3547aa4c66
17 changed files with 156 additions and 82 deletions
--- a/tensorneat/algorithm/neat/gene/node/default.py
+++ b/tensorneat/algorithm/neat/gene/node/default.py
@@ -95,11 +95,17 @@ class DefaultNodeGene(BaseNodeGene):
                (node1[4] != node2[4])
        )

-    def forward(self, attrs, inputs):
+    def forward(self, attrs, inputs, is_output_node=False):
        bias, res, act_idx, agg_idx = attrs

        z = agg(agg_idx, inputs, self.aggregation_options)
        z = bias + res * z
-        z = act(act_idx, z, self.activation_options)
+
+        # the last output node should not be activated
+        z = jax.lax.cond(
+            is_output_node,
+            lambda: z,
+            lambda: act(act_idx, z, self.activation_options)
+        )

        return z