Couple of suggestions for `nn.Sequential`.

Hi @patrick-kidger ,

I was wondering if the following modifications are worth adding to nn.Sequential

1. Supporting OrderedDicts. They have an ordering but also make it easier to access specific indices based on names. Possible bug: At the moment it is possible to only initialize a sequential with an OrderedDict (but it shouldn’t?).

2. Supporting delete, insert, update, append to edit the state of layers. Given that __getitem__ already exists, these would also be qol additions.

Thanks.

[patrick-kidger]

Currently the layers attribute of Sequential can be a list – which of course supports insert, append, ect … However it may make sense to force this to be a tuple to keep it immutable.

Yeah, that’d make sense.

This is pretty common with transfer learning – i.e. doing “model surgery” on a pretrained model to adapt it to a new task. This is where the mutability of PyTorch Modules is really convenient. It would be pretty cool if equinox Modules could be unfrozen and refrozen to allow PyTorch-like manipulation

For performing model surgery - or indeed any kind of PyTree manipulation - I’d recommend eqx.tree_at. For example:

mlp = eqx.nn.MLP(...)
new_linear = eqx.nn.Linear(...)
mlp2 = eqx.tree_at(lambda m: m.layers[-1], mlp, new_linear)

Which replaces the last layer of the MLP.

Does this generally solve your use cases here?

I’ve been thinking of adding an “advanced tricks” section to the docs describing this kind of thing.

[jenkspt]

Is it common to need to change which layers you have in a module after-the-fact?

This is pretty common with transfer learning – i.e. doing “model surgery” on a pretrained model to adapt it to a new task. This is where the mutability of PyTorch Modules is really convenient. It would be pretty cool if equinox Modules could be unfrozen and refrozen to allow PyTorch-like manipulation