use black format all files;

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

tensorneat/examples/brax/ant.py

@@ -16,7 +16,8 @@ if __name__ == '__main__':
                     node_gene=DefaultNodeGene(
                         activation_options=(Act.tanh,),
                         activation_default=Act.tanh,
-                    )
+                    ),
+                    output_transform=Act.tanh
                 ),
                 pop_size=1000,
                 species_size=10,