use black format all files;

remove "return state" for functions which will be executed in vmap;
recover randkey as args in mutation methods

tensorneat/algorithm/neat/genome/base.py

@@ -7,13 +7,13 @@ class BaseGenome:
     network_type = None
 
     def __init__(
-            self,
-            num_inputs: int,
-            num_outputs: int,
-            max_nodes: int,
-            max_conns: int,
-            node_gene: BaseNodeGene = DefaultNodeGene(),
-            conn_gene: BaseConnGene = DefaultConnGene(),
+        self,
+        num_inputs: int,
+        num_outputs: int,
+        max_nodes: int,
+        max_conns: int,
+        node_gene: BaseNodeGene = DefaultNodeGene(),
+        conn_gene: BaseConnGene = DefaultConnGene(),
     ):
         self.num_inputs = num_inputs
         self.num_outputs = num_outputs
@@ -25,6 +25,8 @@ class BaseGenome:
         self.conn_gene = conn_gene
 
     def setup(self, state=State()):
+        state = self.node_gene.setup(state)
+        state = self.conn_gene.setup(state)
         return state
 
     def transform(self, state, nodes, conns):

tensorneat/algorithm/neat/genome/default.py

@@ -1,7 +1,7 @@
 from typing import Callable
 
 import jax, jax.numpy as jnp
-from utils import unflatten_conns, topological_sort, I_INT
+from utils import unflatten_conns, topological_sort, I_INF
 
 from . import BaseGenome
 from ..gene import BaseNodeGene, BaseConnGene, DefaultNodeGene, DefaultConnGene
@@ -10,18 +10,21 @@ from ..gene import BaseNodeGene, BaseConnGene, DefaultNodeGene, DefaultConnGene
 class DefaultGenome(BaseGenome):
     """Default genome class, with the same behavior as the NEAT-Python"""
 
-    network_type = 'feedforward'
+    network_type = "feedforward"
 
-    def __init__(self,
-                 num_inputs: int,
-                 num_outputs: int,
-                 max_nodes=5,
-                 max_conns=4,
-                 node_gene: BaseNodeGene = DefaultNodeGene(),
-                 conn_gene: BaseConnGene = DefaultConnGene(),
-                 output_transform: Callable = None
-                 ):
-        super().__init__(num_inputs, num_outputs, max_nodes, max_conns, node_gene, conn_gene)
+    def __init__(
+        self,
+        num_inputs: int,
+        num_outputs: int,
+        max_nodes=5,
+        max_conns=4,
+        node_gene: BaseNodeGene = DefaultNodeGene(),
+        conn_gene: BaseConnGene = DefaultConnGene(),
+        output_transform: Callable = None,
+    ):
+        super().__init__(
+            num_inputs, num_outputs, max_nodes, max_conns, node_gene, conn_gene
+        )
 
         if output_transform is not None:
             try:
@@ -38,7 +41,7 @@ class DefaultGenome(BaseGenome):
         u_conns = jnp.where(conn_enable, u_conns[1:, :], jnp.nan)
         seqs = topological_sort(nodes, conn_enable)
 
-        return state, seqs, nodes, u_conns
+        return seqs, nodes, u_conns
 
     def forward(self, state, inputs, transformed):
         cal_seqs, nodes, conns = transformed
@@ -49,32 +52,34 @@ class DefaultGenome(BaseGenome):
         nodes_attrs = nodes[:, 1:]
 
         def cond_fun(carry):
-            state_, values, idx = carry
-            return (idx < N) & (cal_seqs[idx] != I_INT)
+            values, idx = carry
+            return (idx < N) & (cal_seqs[idx] != I_INF)
 
         def body_func(carry):
-            state_, values, idx = carry
+            values, idx = carry
             i = cal_seqs[idx]
 
             def hit():
-                s, ins = jax.vmap(self.conn_gene.forward,
-                                  in_axes=(None, 1, 0), out_axes=(None, 0))(state_, conns[:, :, i], values)
-                s, z = self.node_gene.forward(s, nodes_attrs[i], ins, is_output_node=jnp.isin(i, self.output_idx))
+                ins = jax.vmap(self.conn_gene.forward, in_axes=(None, 1, 0))(
+                    state, conns[:, :, i], values
+                )
+                z = self.node_gene.forward(
+                    state,
+                    nodes_attrs[i],
+                    ins,
+                    is_output_node=jnp.isin(i, self.output_idx),
+                )
                 new_values = values.at[i].set(z)
-                return s, new_values
+                return new_values
 
             # the val of input nodes is obtained by the task, not by calculation
-            state_, values = jax.lax.cond(
-                jnp.isin(i, self.input_idx),
-                lambda: (state_, values),
-                hit
-            )
+            values = jax.lax.cond(jnp.isin(i, self.input_idx), lambda: values, hit)
 
-            return state_, values, idx + 1
+            return values, idx + 1
 
-        state, vals, _ = jax.lax.while_loop(cond_fun, body_func, (state, ini_vals, 0))
+        vals, _ = jax.lax.while_loop(cond_fun, body_func, (ini_vals, 0))
 
         if self.output_transform is None:
-            return state, vals[self.output_idx]
+            return vals[self.output_idx]
         else:
-            return state, self.output_transform(vals[self.output_idx])
+            return self.output_transform(vals[self.output_idx])

tensorneat/algorithm/neat/genome/recurrent.py

@@ -10,19 +10,22 @@ from ..gene import BaseNodeGene, BaseConnGene, DefaultNodeGene, DefaultConnGene
 class RecurrentGenome(BaseGenome):
     """Default genome class, with the same behavior as the NEAT-Python"""
 
-    network_type = 'recurrent'
+    network_type = "recurrent"
 
-    def __init__(self,
-                 num_inputs: int,
-                 num_outputs: int,
-                 max_nodes: int,
-                 max_conns: int,
-                 node_gene: BaseNodeGene = DefaultNodeGene(),
-                 conn_gene: BaseConnGene = DefaultConnGene(),
-                 activate_time: int = 10,
-                 output_transform: Callable = None
-                 ):
-        super().__init__(num_inputs, num_outputs, max_nodes, max_conns, node_gene, conn_gene)
+    def __init__(
+        self,
+        num_inputs: int,
+        num_outputs: int,
+        max_nodes: int,
+        max_conns: int,
+        node_gene: BaseNodeGene = DefaultNodeGene(),
+        conn_gene: BaseConnGene = DefaultConnGene(),
+        activate_time: int = 10,
+        output_transform: Callable = None,
+    ):
+        super().__init__(
+            num_inputs, num_outputs, max_nodes, max_conns, node_gene, conn_gene
+        )
         self.activate_time = activate_time
 
         if output_transform is not None:
@@ -39,45 +42,37 @@ class RecurrentGenome(BaseGenome):
         conn_enable = u_conns[0] == 1
         u_conns = jnp.where(conn_enable, u_conns[1:, :], jnp.nan)
 
-        return state, nodes, u_conns
+        return nodes, u_conns
 
     def forward(self, state, inputs, transformed):
         nodes, conns = transformed
 
         N = nodes.shape[0]
         vals = jnp.full((N,), jnp.nan)
-        nodes_attrs = nodes[:, 1:]
+        nodes_attrs = nodes[:, 1:]  # remove index
 
-        def body_func(_, carry):
-            state_, values = carry
+        def body_func(_, values):
 
             # set input values
             values = values.at[self.input_idx].set(inputs)
 
             # calculate connections
-            state_, node_ins = jax.vmap(
-                jax.vmap(
-                    self.conn_gene.forward,
-                    in_axes=(None, 1, None),
-                    out_axes=(None, 0)
-                ),
+            node_ins = jax.vmap(
+                jax.vmap(self.conn_gene.forward, in_axes=(None, 1, None)),
                 in_axes=(None, 1, 0),
-                out_axes=(None, 0)
-            )(state_, conns, values)
+            )(state, conns, values)
 
             # calculate nodes
-            is_output_nodes = jnp.isin(
-                jnp.arange(N),
-                self.output_idx
+            is_output_nodes = jnp.isin(jnp.arange(N), self.output_idx)
+            values = jax.vmap(self.node_gene.forward, in_axes=(None, 0, 0, 0))(
+                state, nodes_attrs, node_ins.T, is_output_nodes
             )
-            state_, values = jax.vmap(
-                self.node_gene.forward,
-                in_axes=(None, 0, 0, 0),
-                out_axes=(None, 0)
-            )(state_, nodes_attrs, node_ins.T, is_output_nodes)
 
-            return state_, values
+            return values
 
-        state, vals = jax.lax.fori_loop(0, self.activate_time, body_func, (state, vals))
+        vals = jax.lax.fori_loop(0, self.activate_time, body_func, vals)
 
-        return state, vals[self.output_idx]
+        if self.output_transform is None:
+            return vals[self.output_idx]
+        else:
+            return self.output_transform(vals[self.output_idx])