diff --git a/examples/tmp.py b/examples/tmp.py
index e900e5a..1d08549 100644
--- a/examples/tmp.py
+++ b/examples/tmp.py
@@ -1,7 +1,7 @@
 import jax, jax.numpy as jnp
 
 from tensorneat.algorithm import NEAT
-from tensorneat.algorithm.neat import DefaultGenome, RecurrentGenome
+from tensorneat.genome import DefaultGenome, RecurrentGenome
 
 key = jax.random.key(0)
 genome = DefaultGenome(num_inputs=5, num_outputs=3, max_nodes=100, max_conns=500, init_hidden_layers=(1, 2 ,3))
diff --git a/network.svg b/network.svg
index 2989755..d86c21f 100644
--- a/network.svg
+++ b/network.svg
@@ -6,7 +6,7 @@
   <rdf:RDF xmlns:dc="http://purl.org/dc/elements/1.1/" xmlns:cc="http://creativecommons.org/ns#" xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#">
    <cc:Work>
     <dc:type rdf:resource="http://purl.org/dc/dcmitype/StillImage"/>
-    <dc:date>2024-07-10T15:27:16.806503</dc:date>
+    <dc:date>2024-07-10T16:50:19.947855</dc:date>
     <dc:format>image/svg+xml</dc:format>
     <dc:creator>
      <cc:Agent>
@@ -32,222 +32,222 @@ z
    <g id="patch_2">
     <path d="M 44.79098 308.403612 
 Q 87.590594 244.204191 129.770035 180.93503 
-" clip-path="url(#p80fa8c6777)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
     <path d="M 125.887134 183.153831 
 L 129.770035 180.93503 
 L 129.215335 185.372632 
 z
-" clip-path="url(#p80fa8c6777)" style="stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="stroke: #000000; stroke-linecap: round"/>
    </g>
    <g id="patch_3">
     <path d="M 46.916335 239.00779 
 Q 87.591722 208.50125 127.372682 178.665529 
-" clip-path="url(#p80fa8c6777)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
     <path d="M 122.972682 179.465529 
 L 127.372682 178.665529 
 L 125.372682 182.665529 
 z
-" clip-path="url(#p80fa8c6777)" style="stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="stroke: #000000; stroke-linecap: round"/>
    </g>
    <g id="patch_4">
     <path d="M 48.647998 172.8 
 Q 87.590519 172.8 125.415005 172.8 
-" clip-path="url(#p80fa8c6777)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
     <path d="M 121.415005 170.8 
 L 125.415005 172.8 
 L 121.415005 174.8 
 z
-" clip-path="url(#p80fa8c6777)" style="stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="stroke: #000000; stroke-linecap: round"/>
    </g>
    <g id="patch_5">
     <path d="M 46.916335 106.59221 
 Q 87.591722 137.09875 127.372682 166.934471 
-" clip-path="url(#p80fa8c6777)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
     <path d="M 125.372682 162.934471 
 L 127.372682 166.934471 
 L 122.972682 166.134471 
 z
-" clip-path="url(#p80fa8c6777)" style="stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="stroke: #000000; stroke-linecap: round"/>
    </g>
    <g id="patch_6">
     <path d="M 44.79098 37.196388 
 Q 87.590594 101.395809 129.770035 164.66497 
-" clip-path="url(#p80fa8c6777)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
     <path d="M 129.215335 160.227368 
 L 129.770035 164.66497 
 L 125.887134 162.446169 
 z
-" clip-path="url(#p80fa8c6777)" style="stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="stroke: #000000; stroke-linecap: round"/>
    </g>
    <g id="patch_7">
     <path d="M 143.30257 175.840943 
 Q 182.796502 190.651168 221.243586 205.068824 
-" clip-path="url(#p80fa8c6777)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
     <path d="M 218.200516 201.791672 
 L 221.243586 205.068824 
 L 216.796023 205.536989 
 z
-" clip-path="url(#p80fa8c6777)" style="stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="stroke: #000000; stroke-linecap: round"/>
    </g>
    <g id="patch_8">
     <path d="M 143.30257 169.759057 
 Q 182.796502 154.948832 221.243586 140.531176 
-" clip-path="url(#p80fa8c6777)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
     <path d="M 216.796023 140.063011 
 L 221.243586 140.531176 
 L 218.200516 143.808328 
 z
-" clip-path="url(#p80fa8c6777)" style="stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="stroke: #000000; stroke-linecap: round"/>
    </g>
    <g id="patch_9">
     <path d="M 238.509181 211.543422 
 Q 278.003113 226.353647 316.450198 240.771303 
-" clip-path="url(#p80fa8c6777)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
     <path d="M 313.407128 237.494151 
 L 316.450198 240.771303 
 L 312.002634 241.239468 
 z
-" clip-path="url(#p80fa8c6777)" style="stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="stroke: #000000; stroke-linecap: round"/>
    </g>
    <g id="patch_10">
     <path d="M 238.509181 205.461536 
 Q 278.003113 190.651312 316.450198 176.233655 
-" clip-path="url(#p80fa8c6777)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
     <path d="M 312.002634 175.765491 
 L 316.450198 176.233655 
 L 313.407128 179.510807 
 z
-" clip-path="url(#p80fa8c6777)" style="stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="stroke: #000000; stroke-linecap: round"/>
    </g>
    <g id="patch_11">
     <path d="M 236.155746 202.027265 
 Q 278.00531 154.946506 319.112092 108.701376 
-" clip-path="url(#p80fa8c6777)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
     <path d="M 314.959818 110.362286 
 L 319.112092 108.701376 
 L 317.949455 113.019741 
 z
-" clip-path="url(#p80fa8c6777)" style="stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="stroke: #000000; stroke-linecap: round"/>
    </g>
    <g id="patch_12">
     <path d="M 236.155746 143.572735 
 Q 278.00531 190.653494 319.112092 236.898624 
-" clip-path="url(#p80fa8c6777)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
     <path d="M 317.949455 232.580259 
 L 319.112092 236.898624 
 L 314.959818 235.237714 
 z
-" clip-path="url(#p80fa8c6777)" style="stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="stroke: #000000; stroke-linecap: round"/>
    </g>
    <g id="patch_13">
     <path d="M 238.509181 140.138464 
 Q 278.003113 154.948688 316.450198 169.366345 
-" clip-path="url(#p80fa8c6777)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
     <path d="M 313.407128 166.089193 
 L 316.450198 169.366345 
 L 312.002634 169.834509 
 z
-" clip-path="url(#p80fa8c6777)" style="stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="stroke: #000000; stroke-linecap: round"/>
    </g>
    <g id="patch_14">
     <path d="M 238.509181 134.056578 
 Q 278.003113 119.246353 316.450198 104.828697 
-" clip-path="url(#p80fa8c6777)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
     <path d="M 312.002634 104.360532 
 L 316.450198 104.828697 
 L 313.407128 108.105849 
 z
-" clip-path="url(#p80fa8c6777)" style="stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="stroke: #000000; stroke-linecap: round"/>
    </g>
    <g id="patch_15">
     <path d="M 334.267833 244.204959 
 Q 373.210353 244.204959 411.03484 244.204959 
-" clip-path="url(#p80fa8c6777)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
     <path d="M 407.03484 242.204959 
 L 411.03484 244.204959 
 L 407.03484 246.204959 
 z
-" clip-path="url(#p80fa8c6777)" style="stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="stroke: #000000; stroke-linecap: round"/>
    </g>
    <g id="patch_16">
     <path d="M 332.53617 239.00779 
 Q 373.211557 208.50125 412.992517 178.665529 
-" clip-path="url(#p80fa8c6777)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
     <path d="M 408.592517 179.465529 
 L 412.992517 178.665529 
 L 410.992517 182.665529 
 z
-" clip-path="url(#p80fa8c6777)" style="stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="stroke: #000000; stroke-linecap: round"/>
    </g>
    <g id="patch_17">
     <path d="M 330.410815 236.998654 
 Q 373.210429 172.799232 415.38987 109.530071 
-" clip-path="url(#p80fa8c6777)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
     <path d="M 411.506968 111.748872 
 L 415.38987 109.530071 
 L 414.83517 113.967673 
 z
-" clip-path="url(#p80fa8c6777)" style="stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="stroke: #000000; stroke-linecap: round"/>
    </g>
    <g id="patch_18">
     <path d="M 332.53617 177.997169 
 Q 373.211557 208.503709 412.992517 238.339429 
-" clip-path="url(#p80fa8c6777)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
     <path d="M 410.992517 234.339429 
 L 412.992517 238.339429 
 L 408.592517 237.539429 
 z
-" clip-path="url(#p80fa8c6777)" style="stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="stroke: #000000; stroke-linecap: round"/>
    </g>
    <g id="patch_19">
     <path d="M 334.267833 172.8 
 Q 373.210353 172.8 411.03484 172.8 
-" clip-path="url(#p80fa8c6777)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
     <path d="M 407.03484 170.8 
 L 411.03484 172.8 
 L 407.03484 174.8 
 z
-" clip-path="url(#p80fa8c6777)" style="stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="stroke: #000000; stroke-linecap: round"/>
    </g>
    <g id="patch_20">
     <path d="M 332.53617 167.602831 
 Q 373.211557 137.096291 412.992517 107.260571 
-" clip-path="url(#p80fa8c6777)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
     <path d="M 408.592517 108.060571 
 L 412.992517 107.260571 
 L 410.992517 111.260571 
 z
-" clip-path="url(#p80fa8c6777)" style="stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="stroke: #000000; stroke-linecap: round"/>
    </g>
    <g id="patch_21">
     <path d="M 330.410815 108.601346 
 Q 373.210429 172.800768 415.38987 236.069929 
-" clip-path="url(#p80fa8c6777)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
     <path d="M 414.83517 231.632327 
 L 415.38987 236.069929 
 L 411.506968 233.851128 
 z
-" clip-path="url(#p80fa8c6777)" style="stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="stroke: #000000; stroke-linecap: round"/>
    </g>
    <g id="patch_22">
     <path d="M 332.53617 106.59221 
 Q 373.211557 137.09875 412.992517 166.934471 
-" clip-path="url(#p80fa8c6777)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
     <path d="M 410.992517 162.934471 
 L 412.992517 166.934471 
 L 408.592517 166.134471 
 z
-" clip-path="url(#p80fa8c6777)" style="stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="stroke: #000000; stroke-linecap: round"/>
    </g>
    <g id="patch_23">
     <path d="M 334.267833 101.395041 
 Q 373.210353 101.395041 411.03484 101.395041 
-" clip-path="url(#p80fa8c6777)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="fill: none; stroke: #000000; stroke-linecap: round"/>
     <path d="M 407.03484 99.395041 
 L 411.03484 101.395041 
 L 407.03484 103.395041 
 z
-" clip-path="url(#p80fa8c6777)" style="stroke: #000000; stroke-linecap: round"/>
+" clip-path="url(#p8fe09283f8)" style="stroke: #000000; stroke-linecap: round"/>
    </g>
    <g id="PathCollection_1">
     <path d="M 39.986777 324.270171 
@@ -260,7 +260,7 @@ C 32.239022 311.110224 31.326523 313.313191 31.326523 315.609917
 C 31.326523 317.906644 32.239022 320.109611 33.863053 321.733642 
 C 35.487083 323.357672 37.690051 324.270171 39.986777 324.270171 
 z
-" clip-path="url(#p80fa8c6777)" style="fill: #0000ff; stroke: #0000ff"/>
+" clip-path="url(#p8fe09283f8)" style="fill: #0000ff; stroke: #0000ff"/>
     <path d="M 39.986777 252.865213 
 C 42.283503 252.865213 44.486471 251.952714 46.110501 250.328683 
 C 47.734532 248.704652 48.647031 246.501685 48.647031 244.204959 
@@ -271,7 +271,7 @@ C 32.239022 239.705265 31.326523 241.908232 31.326523 244.204959
 C 31.326523 246.501685 32.239022 248.704652 33.863053 250.328683 
 C 35.487083 251.952714 37.690051 252.865213 39.986777 252.865213 
 z
-" clip-path="url(#p80fa8c6777)" style="fill: #0000ff; stroke: #0000ff"/>
+" clip-path="url(#p8fe09283f8)" style="fill: #0000ff; stroke: #0000ff"/>
     <path d="M 39.986777 181.460254 
 C 42.283503 181.460254 44.486471 180.547755 46.110501 178.923724 
 C 47.734532 177.299694 48.647031 175.096726 48.647031 172.8 
@@ -282,7 +282,7 @@ C 32.239022 168.300306 31.326523 170.503274 31.326523 172.8
 C 31.326523 175.096726 32.239022 177.299694 33.863053 178.923724 
 C 35.487083 180.547755 37.690051 181.460254 39.986777 181.460254 
 z
-" clip-path="url(#p80fa8c6777)" style="fill: #0000ff; stroke: #0000ff"/>
+" clip-path="url(#p8fe09283f8)" style="fill: #0000ff; stroke: #0000ff"/>
     <path d="M 39.986777 110.055295 
 C 42.283503 110.055295 44.486471 109.142796 46.110501 107.518766 
 C 47.734532 105.894735 48.647031 103.691768 48.647031 101.395041 
@@ -293,7 +293,7 @@ C 32.239022 96.895348 31.326523 99.098315 31.326523 101.395041
 C 31.326523 103.691768 32.239022 105.894735 33.863053 107.518766 
 C 35.487083 109.142796 37.690051 110.055295 39.986777 110.055295 
 z
-" clip-path="url(#p80fa8c6777)" style="fill: #0000ff; stroke: #0000ff"/>
+" clip-path="url(#p8fe09283f8)" style="fill: #0000ff; stroke: #0000ff"/>
     <path d="M 39.986777 38.650337 
 C 42.283503 38.650337 44.486471 37.737838 46.110501 36.113807 
 C 47.734532 34.489776 48.647031 32.286809 48.647031 29.990083 
@@ -304,7 +304,7 @@ C 32.239022 25.490389 31.326523 27.693356 31.326523 29.990083
 C 31.326523 32.286809 32.239022 34.489776 33.863053 36.113807 
 C 35.487083 37.737838 37.690051 38.650337 39.986777 38.650337 
 z
-" clip-path="url(#p80fa8c6777)" style="fill: #0000ff; stroke: #0000ff"/>
+" clip-path="url(#p8fe09283f8)" style="fill: #0000ff; stroke: #0000ff"/>
     <path d="M 135.193388 181.460254 
 C 137.490115 181.460254 139.693082 180.547755 141.317113 178.923724 
 C 142.941143 177.299694 143.853642 175.096726 143.853642 172.8 
@@ -315,7 +315,7 @@ C 127.445633 168.300306 126.533134 170.503274 126.533134 172.8
 C 126.533134 175.096726 127.445633 177.299694 129.069664 178.923724 
 C 130.693695 180.547755 132.896662 181.460254 135.193388 181.460254 
 z
-" clip-path="url(#p80fa8c6777)" style="fill: #0000ff; stroke: #0000ff"/>
+" clip-path="url(#p8fe09283f8)" style="fill: #0000ff; stroke: #0000ff"/>
     <path d="M 230.4 217.162733 
 C 232.696726 217.162733 234.899694 216.250234 236.523724 214.626204 
 C 238.147755 213.002173 239.060254 210.799206 239.060254 208.502479 
@@ -326,7 +326,7 @@ C 222.652245 204.002786 221.739746 206.205753 221.739746 208.502479
 C 221.739746 210.799206 222.652245 213.002173 224.276276 214.626204 
 C 225.900306 216.250234 228.103274 217.162733 230.4 217.162733 
 z
-" clip-path="url(#p80fa8c6777)" style="fill: #0000ff; stroke: #0000ff"/>
+" clip-path="url(#p8fe09283f8)" style="fill: #0000ff; stroke: #0000ff"/>
     <path d="M 230.4 145.757775 
 C 232.696726 145.757775 234.899694 144.845276 236.523724 143.221245 
 C 238.147755 141.597214 239.060254 139.394247 239.060254 137.097521 
@@ -337,7 +337,7 @@ C 222.652245 132.597827 221.739746 134.800794 221.739746 137.097521
 C 221.739746 139.394247 222.652245 141.597214 224.276276 143.221245 
 C 225.900306 144.845276 228.103274 145.757775 230.4 145.757775 
 z
-" clip-path="url(#p80fa8c6777)" style="fill: #0000ff; stroke: #0000ff"/>
+" clip-path="url(#p8fe09283f8)" style="fill: #0000ff; stroke: #0000ff"/>
     <path d="M 325.606612 252.865213 
 C 327.903338 252.865213 330.106305 251.952714 331.730336 250.328683 
 C 333.354367 248.704652 334.266866 246.501685 334.266866 244.204959 
@@ -348,7 +348,7 @@ C 317.858857 239.705265 316.946358 241.908232 316.946358 244.204959
 C 316.946358 246.501685 317.858857 248.704652 319.482887 250.328683 
 C 321.106918 251.952714 323.309885 252.865213 325.606612 252.865213 
 z
-" clip-path="url(#p80fa8c6777)" style="fill: #0000ff; stroke: #0000ff"/>
+" clip-path="url(#p8fe09283f8)" style="fill: #0000ff; stroke: #0000ff"/>
     <path d="M 325.606612 181.460254 
 C 327.903338 181.460254 330.106305 180.547755 331.730336 178.923724 
 C 333.354367 177.299694 334.266866 175.096726 334.266866 172.8 
@@ -359,7 +359,7 @@ C 317.858857 168.300306 316.946358 170.503274 316.946358 172.8
 C 316.946358 175.096726 317.858857 177.299694 319.482887 178.923724 
 C 321.106918 180.547755 323.309885 181.460254 325.606612 181.460254 
 z
-" clip-path="url(#p80fa8c6777)" style="fill: #0000ff; stroke: #0000ff"/>
+" clip-path="url(#p8fe09283f8)" style="fill: #0000ff; stroke: #0000ff"/>
     <path d="M 325.606612 110.055295 
 C 327.903338 110.055295 330.106305 109.142796 331.730336 107.518766 
 C 333.354367 105.894735 334.266866 103.691768 334.266866 101.395041 
@@ -370,7 +370,7 @@ C 317.858857 96.895348 316.946358 99.098315 316.946358 101.395041
 C 316.946358 103.691768 317.858857 105.894735 319.482887 107.518766 
 C 321.106918 109.142796 323.309885 110.055295 325.606612 110.055295 
 z
-" clip-path="url(#p80fa8c6777)" style="fill: #0000ff; stroke: #0000ff"/>
+" clip-path="url(#p8fe09283f8)" style="fill: #0000ff; stroke: #0000ff"/>
     <path d="M 420.813223 252.865213 
 C 423.109949 252.865213 425.312917 251.952714 426.936947 250.328683 
 C 428.560978 248.704652 429.473477 246.501685 429.473477 244.204959 
@@ -381,7 +381,7 @@ C 413.065468 239.705265 412.152969 241.908232 412.152969 244.204959
 C 412.152969 246.501685 413.065468 248.704652 414.689499 250.328683 
 C 416.313529 251.952714 418.516497 252.865213 420.813223 252.865213 
 z
-" clip-path="url(#p80fa8c6777)" style="fill: #0000ff; stroke: #0000ff"/>
+" clip-path="url(#p8fe09283f8)" style="fill: #0000ff; stroke: #0000ff"/>
     <path d="M 420.813223 181.460254 
 C 423.109949 181.460254 425.312917 180.547755 426.936947 178.923724 
 C 428.560978 177.299694 429.473477 175.096726 429.473477 172.8 
@@ -392,7 +392,7 @@ C 413.065468 168.300306 412.152969 170.503274 412.152969 172.8
 C 412.152969 175.096726 413.065468 177.299694 414.689499 178.923724 
 C 416.313529 180.547755 418.516497 181.460254 420.813223 181.460254 
 z
-" clip-path="url(#p80fa8c6777)" style="fill: #0000ff; stroke: #0000ff"/>
+" clip-path="url(#p8fe09283f8)" style="fill: #0000ff; stroke: #0000ff"/>
     <path d="M 420.813223 110.055295 
 C 423.109949 110.055295 425.312917 109.142796 426.936947 107.518766 
 C 428.560978 105.894735 429.473477 103.691768 429.473477 101.395041 
@@ -403,12 +403,12 @@ C 413.065468 96.895348 412.152969 99.098315 412.152969 101.395041
 C 412.152969 103.691768 413.065468 105.894735 414.689499 107.518766 
 C 416.313529 109.142796 418.516497 110.055295 420.813223 110.055295 
 z
-" clip-path="url(#p80fa8c6777)" style="fill: #0000ff; stroke: #0000ff"/>
+" clip-path="url(#p8fe09283f8)" style="fill: #0000ff; stroke: #0000ff"/>
    </g>
   </g>
  </g>
  <defs>
-  <clipPath id="p80fa8c6777">
+  <clipPath id="p8fe09283f8">
    <rect x="0" y="0" width="460.8" height="345.6"/>
   </clipPath>
  </defs>
diff --git a/tensorneat/algorithm/neat/__init__.py b/tensorneat/algorithm/neat/__init__.py
index 06f14d2..cab968b 100644
--- a/tensorneat/algorithm/neat/__init__.py
+++ b/tensorneat/algorithm/neat/__init__.py
@@ -1,4 +1,2 @@
-from .gene import *
-from .genome import *
 from .species import *
 from .neat import NEAT
diff --git a/tensorneat/algorithm/neat/species/base.py b/tensorneat/algorithm/neat/species/base.py
index f6175de..cf03d4f 100644
--- a/tensorneat/algorithm/neat/species/base.py
+++ b/tensorneat/algorithm/neat/species/base.py
@@ -1,5 +1,5 @@
 from tensorneat.common import State, StatefulBaseClass
-from ..genome import BaseGenome
+from tensorneat.genome import BaseGenome
 
 
 class BaseSpecies(StatefulBaseClass):
diff --git a/tensorneat/algorithm/neat/species/default.py b/tensorneat/algorithm/neat/species/default.py
index 310e070..3c5e82a 100644
--- a/tensorneat/algorithm/neat/species/default.py
+++ b/tensorneat/algorithm/neat/species/default.py
@@ -1,16 +1,17 @@
 import jax, jax.numpy as jnp
+
+from .base import BaseSpecies
 from tensorneat.common import (
     State,
     rank_elements,
     argmin_with_mask,
     fetch_first,
 )
-from ..genome.utils import (
+from tensorneat.genome.utils import (
     extract_conn_attrs,
     extract_node_attrs,
 )
-from ..genome import BaseGenome
-from .base import BaseSpecies
+from tensorneat.genome import BaseGenome
 
 
 """
diff --git a/tensorneat/genome/base.py b/tensorneat/genome/base.py
index 502d3c6..f3091f2 100644
--- a/tensorneat/genome/base.py
+++ b/tensorneat/genome/base.py
@@ -3,7 +3,7 @@ from typing import Callable, Sequence
 import numpy as np
 import jax
 from jax import vmap, numpy as jnp
-from ..gene import BaseNodeGene, BaseConnGene
+from .gene import BaseNodeGene, BaseConnGene
 from .operations import BaseMutation, BaseCrossover, BaseDistance
 from tensorneat.common import (
     State,
diff --git a/tensorneat/genome/default.py b/tensorneat/genome/default.py
index de29eec..73fd736 100644
--- a/tensorneat/genome/default.py
+++ b/tensorneat/genome/default.py
@@ -5,8 +5,8 @@ from jax import vmap, numpy as jnp
 import numpy as np
 import sympy as sp
 
-from . import BaseGenome
-from ..gene import DefaultNodeGene, DefaultConnGene
+from .base import BaseGenome
+from .gene import DefaultNodeGene, DefaultConnGene
 from .operations import DefaultMutation, DefaultCrossover, DefaultDistance
 from .utils import unflatten_conns, extract_node_attrs, extract_conn_attrs
 
@@ -102,7 +102,7 @@ class DefaultGenome(BaseGenome):
                     state,
                     nodes_attrs[i],
                     ins,
-                    is_output_node=jnp.isin(nodes[0], self.output_idx),  # nodes[0] -> the key of nodes
+                    is_output_node=jnp.isin(nodes[i, 0], self.output_idx),  # nodes[0] -> the key of nodes
                 )
 
                 # set new value
diff --git a/tensorneat/genome/recurrent.py b/tensorneat/genome/recurrent.py
index a289e8f..d833265 100644
--- a/tensorneat/genome/recurrent.py
+++ b/tensorneat/genome/recurrent.py
@@ -3,9 +3,9 @@ from jax import vmap, numpy as jnp
 from .utils import unflatten_conns
 
 from .base import BaseGenome
+from .gene import DefaultNodeGene, DefaultConnGene
 from .operations import DefaultMutation, DefaultCrossover, DefaultDistance
 from .utils import unflatten_conns, extract_node_attrs, extract_conn_attrs
-from ..gene import DefaultNodeGene, DefaultConnGene
 
 from tensorneat.common import attach_with_inf