remove attr enable for conn

tensorneat/algorithm/neat/genome/base.py

@@ -106,21 +106,19 @@ class BaseGenome:
             self.input_idx, jnp.full_like(self.input_idx, new_node_key)
         ]
         conns = conns.at[self.input_idx, :2].set(input_conns)  # in-keys, out-keys
-        conns = conns.at[self.input_idx, 2].set(True)  # enable
 
         # output-hidden connections
         output_conns = jnp.c_[
             jnp.full_like(self.output_idx, new_node_key), self.output_idx
         ]
         conns = conns.at[self.output_idx, :2].set(output_conns)  # in-keys, out-keys
-        conns = conns.at[self.output_idx, 2].set(True)  # enable
 
         conn_keys = jax.random.split(k2, num=len(self.input_idx) + len(self.output_idx))
         # generate random attributes for conns
         random_conn_attrs = jax.vmap(
             self.conn_gene.new_random_attrs, in_axes=(None, 0)
         )(state, conn_keys)
-        conns = conns.at[: len(conn_keys), 3:].set(random_conn_attrs)
+        conns = conns.at[: len(conn_keys), 2:].set(random_conn_attrs)
 
         return nodes, conns
 

tensorneat/algorithm/neat/genome/default.py

@@ -45,19 +45,15 @@ class DefaultGenome(BaseGenome):
 
     def transform(self, state, nodes, conns):
         u_conns = unflatten_conns(nodes, conns)
-        conn_enable = u_conns[0] == 1
+        conn_exist = ~jnp.isnan(u_conns[0])
 
-        # remove enable attr
-        u_conns = jnp.where(conn_enable, u_conns[1:, :], jnp.nan)
-        seqs = topological_sort(nodes, conn_enable)
+        seqs = topological_sort(nodes, conn_exist)
 
         return seqs, nodes, u_conns
 
     def restore(self, state, transformed):
         seqs, nodes, u_conns = transformed
         conns = flatten_conns(nodes, u_conns, C=self.max_conns)
-        # restore enable
-        conns = jnp.insert(conns, obj=2, values=1, axis=1)
         return nodes, conns
 
     def forward(self, state, inputs, transformed):
@@ -79,14 +75,15 @@ class DefaultGenome(BaseGenome):
                 ins = jax.vmap(self.conn_gene.forward, in_axes=(None, 1, 0))(
                     state, u_conns[:, :, i], values
                 )
+
                 z = self.node_gene.forward(
                     state,
                     nodes_attrs[i],
                     ins,
                     is_output_node=jnp.isin(i, self.output_idx),
                 )
-                new_values = values.at[i].set(z)
 
+                new_values = values.at[i].set(z)
                 return new_values
 
             # the val of input nodes is obtained by the task, not by calculation

tensorneat/algorithm/neat/genome/recurrent.py

@@ -47,19 +47,11 @@ class RecurrentGenome(BaseGenome):
 
     def transform(self, state, nodes, conns):
         u_conns = unflatten_conns(nodes, conns)
-
-        # remove un-enable connections and remove enable attr
-        conn_enable = u_conns[0] == 1
-        u_conns = jnp.where(conn_enable, u_conns[1:, :], jnp.nan)
-
         return nodes, u_conns
 
     def restore(self, state, transformed):
         nodes, u_conns = transformed
         conns = flatten_conns(nodes, u_conns, C=self.max_conns)
-
-        # restore enable
-        conns = jnp.insert(conns, obj=2, values=1, axis=1)
         return nodes, conns
 
     def forward(self, state, inputs, transformed):