add save and load function for classes.

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