make fully stateful in module ga.

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
+
+    def setup(self, state=State()):
+        return state
+
+    def __call__(self, state, key, genome, nodes, conns, new_node_key):
+        raise NotImplementedError

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)