added docstrings to add_connection

README.md

@@ -25,11 +25,7 @@ Mutations:
   * [x] Weight shift
   * [x] Weight reset
   * [x] Add node
-  * [ ] Add connection
-    * [x] Add a `depth` field to `ActivationState` to track node depths
-    * [x] Test `get_depth` on multiple topologies
-    * [x] Update `depth_outdated` on mutations affecting network topology
-    * [ ] Add tests
+  * [x] Add connection
 * [ ] Wrap all mutations in a single function
 
 Crossing:

neat_jax/mutations.py

@@ -236,6 +236,20 @@ def add_connection(
         max_nodes: int,
         scale_weights: float = 0.1,
     ) -> tuple[Network, ActivationState]:
+        """
+        Samples a new connection between nodes, ensuring it adheres to network topology order
+        (i.e. the depth of the sender is lower than the depth of the receiver).
+
+        Args:
+            key (jax.random.PRNGKey): A PRNG key used for random operations.
+            net (Network): The current state of the network, containing node types and connections.
+            activation_state (ActivationState): The current activation state of the network, including node depths.
+            max_nodes (int): The maximum number of nodes in the network, used for defining array sizes.
+
+        Returns:
+            Network: The mutated network with a new connection
+            ActivationState: The activation state of the network with updated node depths
+        """
 
         def _mutate_fn(
             key: random.PRNGKey,
@@ -356,4 +370,4 @@ def _bypass_apply_mutation(*args):
             operand=None,
         )
 
-        return net
+        return net, activation_state

tests/test_mutations.py

@@ -82,6 +82,9 @@ class MutationTests(chex.TestCase, parameterized.TestCase):
                 "added_connection_node_types": jnp.array(
                     [0, 0, 0, 2, 1, 3, 3, 3, 3, 3]
                 ),
+                "outdated_depths": jnp.bool_(
+                    True
+                ),  # no mutation => depths were not updated
             },
         ),
         (
@@ -130,6 +133,7 @@ class MutationTests(chex.TestCase, parameterized.TestCase):
                 "added_connection_node_types": jnp.array(
                     [0, 0, 0, 2, 1, 3, 3, 3, 3, 3]
                 ),
+                "outdated_depths": jnp.bool_(False),
             },
         ),
         (
@@ -178,6 +182,7 @@ class MutationTests(chex.TestCase, parameterized.TestCase):
                 "added_connection_node_types": jnp.array(
                     [0, 0, 0, 2, 1, 3, 3, 3, 3, 3]
                 ),
+                "outdated_depths": jnp.bool_(False),
             },
         ),
     )
@@ -191,7 +196,7 @@ def test_mutate(
         shifted_weights = self.variant(mutations.weight_shift)(key, net, 0.1).weights
         mutated_weights = self.variant(mutations.weight_mutation)(key, net, 0.1).weights
         added_node_network = self.variant(mutations.add_node)(key, net, 0.1)
-        added_connection_network = self.variant(
+        added_connection_network, added_connection_activation_state = self.variant(
             mutations.add_connection, static_argnames=["self", "max_nodes"]
         )(key, net, activation_state, t_params["max_nodes"])
 
@@ -225,3 +230,7 @@ def test_mutate(
         chex.assert_trees_all_equal(
             added_connection_network.node_types, expected["added_connection_node_types"]
         )
+        assert (
+            added_connection_activation_state.outdated_depths
+            == expected["outdated_depths"]
+        )