Aditya Singh

So what I meant with `prefix` loading currently works (for same save and load type) was this: I modified the test from https://github.com/patrick-kidger/equinox/blob/a89d5b486d13588caffc095f172a2ec39fd68278/tests/test_serialisation.py#L40-L49 to ```python like = ( like_jax_array1, like_jax_array2,...

Thanks for the info on the prefix! I'll try to dig in deeper as to why this case worked for deserialisation.

Hi, As per the implementation of `tree_map`: https://jax.readthedocs.io/en/latest/_modules/jax/_src/tree_util.html#tree_map It also uses `tree_flatten` on the `tree` and on the `rest` which is later combined with `unflatten`. How does `ordered...` differ from...

Hi, Wanted to double check the behaviour of `StateIndex` within `ordered_tree_map` or current `tree_map`. ```python state_1 = eqx.experimental.StateIndex() state_2 = eqx.experimental.StateIndex() eqx.experimental.set_state(state_1, jnp.asarray(4)) eqx.experimental.set_state(state_2, jnp.asarray(2)) z = jtu.tree_map(lambda *xs: tuple(xs),...

Thanks for confirming this.

Unsurprisingly, `ordered_tree_map` fits in well. I wasn't sure if this should be made public or not. At the moment, it is a public method. As per previous discussions [here](https://github.com/patrick-kidger/equinox/issues/136#issuecomment-1196469019) >...

Sorry I did not understand it completely. So `model.eqx` is the serialised complete model. `model.final_linear` will be saved towards the end of the serialised file assuming it happens to be...

Ahh, I see. I think your proposed approach is more elegant. Sorry to stretch this discussion so much. Hopefully, this is the last bit in this thread: I don't think...

Thanks, - The issue persists even after moving to highest precision. - I have opened an issue: https://github.com/google/jax/issues/11790 to follow this up. Feel free to close this issue.

Adding my experience to this issue. Using an IDE (`pycharm`) I am able to get suggestions for the input parameters for most of the subclasses. But, `nn.Sequential` fails to parse...