add save and load function for classes.

tensorneat/algorithm/neat/ga/crossover/base.py

@@ -1,9 +1,6 @@
-from utils import State
+from utils import StatefulBaseClass
 
 
-class BaseCrossover:
-    def setup(self, state=State()):
-        return state
-
+class BaseCrossover(StatefulBaseClass):
     def __call__(self, state, randkey, genome, nodes1, nodes2, conns1, conns2):
         raise NotImplementedError

tensorneat/algorithm/neat/ga/mutation/base.py

@@ -1,9 +1,6 @@
-from utils import State
+from utils import StatefulBaseClass
 
 
-class BaseMutation:
-    def setup(self, state=State()):
-        return state
-
+class BaseMutation(StatefulBaseClass):
     def __call__(self, state, randkey, genome, nodes, conns, new_node_key):
         raise NotImplementedError

tensorneat/algorithm/neat/gene/base.py

@@ -1,8 +1,8 @@
 import jax, jax.numpy as jnp
-from utils import State
+from utils import State, StatefulBaseClass
 
 
-class BaseGene:
+class BaseGene(StatefulBaseClass):
     "Base class for node genes or connection genes."
     fixed_attrs = []
     custom_attrs = []
@@ -10,9 +10,6 @@ class BaseGene:
     def __init__(self):
         pass
 
-    def setup(self, state=State()):
-        return state
-
     def new_identity_attrs(self, state):
         # the attrs which do identity transformation, used in mutate add node
         raise NotImplementedError

tensorneat/algorithm/neat/gene/conn/base.py

@@ -1,4 +1,4 @@
-import jax, jax.numpy as jnp
+import jax
 from .. import BaseGene
 
 

tensorneat/algorithm/neat/genome/base.py

@@ -1,10 +1,10 @@
 import jax, jax.numpy as jnp
 from ..gene import BaseNodeGene, BaseConnGene
 from ..ga import BaseMutation, BaseCrossover
-from utils import State
+from utils import State, StatefulBaseClass
 
 
-class BaseGenome:
+class BaseGenome(StatefulBaseClass):
     network_type = None
 
     def __init__(

tensorneat/algorithm/neat/species/base.py

@@ -1,15 +1,12 @@
-from utils import State
+from utils import State, StatefulBaseClass
 from ..genome import BaseGenome
 
 
-class BaseSpecies:
+class BaseSpecies(StatefulBaseClass):
     genome: BaseGenome
     pop_size: int
     species_size: int
 
-    def setup(self, state=State()):
-        return state
-
     def ask(self, state: State):
         raise NotImplementedError