diff --git a/tensorneat/algorithm/neat/ga/crossover/base.py b/tensorneat/algorithm/neat/ga/crossover/base.py
index 9f84ac9..c7244d3 100644
--- a/tensorneat/algorithm/neat/ga/crossover/base.py
+++ b/tensorneat/algorithm/neat/ga/crossover/base.py
@@ -1,3 +1,10 @@
+from utils import State
+
+
 class BaseCrossover:
-    def __call__(self, randkey, genome, nodes1, nodes2, conns1, conns2):
+
+    def setup(self, state=State()):
+        return state
+
+    def __call__(self, state, key, genome, nodes1, nodes2, conns1, conns2):
         raise NotImplementedError
diff --git a/tensorneat/algorithm/neat/ga/crossover/default.py b/tensorneat/algorithm/neat/ga/crossover/default.py
index deaf745..c7eb531 100644
--- a/tensorneat/algorithm/neat/ga/crossover/default.py
+++ b/tensorneat/algorithm/neat/ga/crossover/default.py
@@ -5,12 +5,12 @@ from .base import BaseCrossover
 
 class DefaultCrossover(BaseCrossover):
 
-    def __call__(self, randkey, genome, nodes1, conns1, nodes2, conns2):
+    def __call__(self, state, key, genome, nodes1, conns1, nodes2, conns2):
         """
         use genome1 and genome2 to generate a new genome
         notice that genome1 should have higher fitness than genome2 (genome1 is winner!)
         """
-        randkey_1, randkey_2, key = jax.random.split(randkey, 3)
+        randkey_1, randkey_2, key = jax.random.split(key, 3)
 
         # crossover nodes
         keys1, keys2 = nodes1[:, 0], nodes2[:, 0]
diff --git a/tensorneat/algorithm/neat/ga/mutation/base.py b/tensorneat/algorithm/neat/ga/mutation/base.py
index 8291522..e56e8e6 100644
--- a/tensorneat/algorithm/neat/ga/mutation/base.py
+++ b/tensorneat/algorithm/neat/ga/mutation/base.py
@@ -1,3 +1,10 @@
+from utils import State
+
+
 class BaseMutation:
-    def __call__(self, key, genome, nodes, conns, new_node_key):
-        raise NotImplementedError
\ No newline at end of file
+
+    def setup(self, state=State()):
+        return state
+
+    def __call__(self, state, key, genome, nodes, conns, new_node_key):
+        raise NotImplementedError
diff --git a/tensorneat/algorithm/neat/ga/mutation/default.py b/tensorneat/algorithm/neat/ga/mutation/default.py
index 331dbf3..2fe4958 100644
--- a/tensorneat/algorithm/neat/ga/mutation/default.py
+++ b/tensorneat/algorithm/neat/ga/mutation/default.py
@@ -17,7 +17,7 @@ class DefaultMutation(BaseMutation):
         self.node_add = node_add
         self.node_delete = node_delete
 
-    def __call__(self, randkey, genome, nodes, conns, new_node_key):
+    def __call__(self, state, randkey, genome, nodes, conns, new_node_key):
         k1, k2 = jax.random.split(randkey)
 
         nodes, conns = self.mutate_structure(k1, genome, nodes, conns, new_node_key)
diff --git a/tensorneat/algorithm/neat/genome/base.py b/tensorneat/algorithm/neat/genome/base.py
index b704f45..fd711ab 100644
--- a/tensorneat/algorithm/neat/genome/base.py
+++ b/tensorneat/algorithm/neat/genome/base.py
@@ -33,7 +33,7 @@ class BaseGenome:
     def forward(self, state, inputs, transformed):
         raise NotImplementedError
 
-    def add_node(self, state, nodes, new_key: int, attrs):
+    def add_node(self, nodes, new_key: int, attrs):
         """
         Add a new node to the genome.
         The new node will place at the first NaN row.
@@ -43,14 +43,14 @@ class BaseGenome:
         new_nodes = nodes.at[pos, 0].set(new_key)
         return new_nodes.at[pos, 1:].set(attrs)
 
-    def delete_node_by_pos(self, state, nodes, pos):
+    def delete_node_by_pos(self, nodes, pos):
         """
         Delete a node from the genome.
         Delete the node by its pos in nodes.
         """
         return nodes.at[pos].set(jnp.nan)
 
-    def add_conn(self, state, conns, i_key, o_key, enable: bool, attrs):
+    def add_conn(self, conns, i_key, o_key, enable: bool, attrs):
         """
         Add a new connection to the genome.
         The new connection will place at the first NaN row.
@@ -60,7 +60,7 @@ class BaseGenome:
         new_conns = conns.at[pos, 0:3].set(jnp.array([i_key, o_key, enable]))
         return new_conns.at[pos, 3:].set(attrs)
 
-    def delete_conn_by_pos(self, state, conns, pos):
+    def delete_conn_by_pos(self, conns, pos):
         """
         Delete a connection from the genome.
         Delete the connection by its idx.