5 fix add node mutation bug

[RPegoud]

Corrected the initialization of network input activations at initialization and inference time:

def get_initial_activations(inputs: chex.Array, senders: chex.Array) -> chex.Array:
    """Initializes ActivationState.values based on an array of inputs."""
    default_values = jnp.zeros_like(senders)
    keys = jnp.arange(len(inputs))
    matches = jnp.searchsorted(keys, senders)
    mask = (matches < len(keys)) & (keys[matches] == senders)
    return jnp.where(mask, inputs[matches], default_values)


def create_toggled_mask(
    node_types: chex.Array, senders: chex.Array, input_size: int
) -> chex.Array:
    """Initializes the ActivationState.toggled mask based on input nodes"""
    input_nodes_indices = jnp.where(node_types == 0, size=input_size)[0]
    sender_mask = jnp.isin(senders, input_nodes_indices)
    return jnp.int32(sender_mask)


def activation_state_from_inputs(
    inputs: chex.Array,
    senders: chex.Array,
    node_types: chex.Array,
    input_size: int,
    max_nodes: int,
) -> "ActivationState":
    """
    Resets the ActivationState in prevision of a forward pass or a depth scan.

    Args:
        inputs (chex.Array): The activation values of the network's input nodes
        max_nodes (int): The maximum capacity of the network

    Returns:
        ActivationState: The reset ActivationState with:

            - ``values``: initialized based on inputs
            - ``toggled``: input neurons toggled
            - ``activation_counts``: set to zero
            - ``has_fired``: set to zero
            - ``outdated_depths``: True
    """

    values = get_initial_activations(inputs, senders)
    toggled = create_toggled_mask(node_types, senders, input_size)

    return ActivationState(
        values=values,
        toggled=toggled,
        activation_counts=jnp.zeros(max_nodes, dtype=jnp.int32),
        has_fired=jnp.zeros(max_nodes, dtype=jnp.int32),
        node_depths=jnp.zeros(max_nodes, dtype=jnp.int32),
        outdated_depths=True,
    )