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/__init__.py

@@ -1,2 +1,2 @@
 from .base import BaseAlgorithm
-from .neat import NEAT
+from .neat import NEAT