Move more things to pytrees, at least internally

[jnsbck]

I think we should think about moving more and more things to pytrees, i.e. recordings, externals. This would enable us to use tree_map etc. to operate on them more efficiently. Related to #555.

For example, we can make channels and synapses pytrees. This way, we would not need to add the params to jaxnodes / jaxedges seperately. This also means params / states could be nested, i.e. useful for NODEs as well.

class Mechanism:
    def __init__(self):
        self.params = {"kp1": 1, "p2": 2}
        self.states = {"s1": 1, "s2": 2}
        self.name = "mech_name"

    def update_states(*args, **kwargs):
        return

    def tree_flatten(self):
        return ((self.params, self.states), self.name)

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        obj = cls.__new__(cls)
        obj.params, obj.states = children
        obj.name = aux_data
        return obj

tree_util.register_pytree_node(
    Mechanism,
    Mechanism.tree_flatten,
    Mechanism.tree_unflatten
)

The Modules themselves could also be potentially converted to pytrees.