add "update_by_batch" in gene;

add flatten_conns as an inverse function for unflatten_conns;
add "test_flatten.ipynb" as test for them.

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

@@ -29,3 +29,10 @@ class BaseConnGene(BaseGene):
 
     def forward(self, state, attrs, inputs):
         raise NotImplementedError
+
+    def update_by_batch(self, state, attrs, batch_inputs):
+        # default: do not update attrs, but to calculate batch_res
+        return (
+            jax.vmap(self.forward, in_axes=(None, None, 0))(state, attrs, batch_inputs),
+            attrs,
+        )