Make build_dynamic_state_utils compatible with networks

[Matthijspals]

We should not filter NaNs and branch points out of jaxnodes.edges / synapses:

but just use tree_map helper functions that only apply the filtering functions fn to leaves with keys in valid_keys (those that are in jaxnodes.keys())

def tree_map_leaves_with_valid_key(tree, fn, valid_keys=None):
    """
    Apply fn(leaf) selectively to leaves that satisfy is_valid_membrane_leaf(key, leaf, valid_keys).
    """
    return tree_map_with_path(
        lambda path, leaf: (
            fn(leaf)
            if is_valid_membrane_leaf(get_key_name(path), leaf, valid_keys)
            else leaf
        ),
        tree,
    )