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/species/base.py

@@ -1,15 +1,20 @@
 from utils import State
+from ..genome import BaseGenome
 
 
 class BaseSpecies:
-    def setup(self, key, state=State()):
+    genome: BaseGenome
+    pop_size: int
+    species_size: int
+
+    def setup(self, state=State()):
         return state
 
     def ask(self, state: State):
         raise NotImplementedError
 
-    def update_species(self, state, fitness, generation):
+    def update_species(self, state, fitness):
         raise NotImplementedError
 
-    def speciate(self, state, generation):
+    def speciate(self, state):
         raise NotImplementedError