Better compatibility between eqx.filter_eval_shape and eqx.tree_serialise_leaves

[patrick-kidger]

Right now it's possible to load a model from a checkpoint by doing something like:

def run(..., checkpoint=None):
    model = Model(...)
    if checkpoint is not None:
        model = eqx.tree_deserialise_leaves(checkpoint, model)
    ...

If one really cares about efficiency then the first line can be replaced by model = eqx.filter_eval_shape(Model, ...). However this returns ShapeDtypeStructs rather than arrays, so the deserialisation line won't work. Right now you need to manually hack something together. We should instead make this easier to do.