Using Frozen Parameters to Freeze Differing Parameters Based Upon Epoch

[adam-hartshorne]

I am trying to figure out if it is possible to efficiently use something like the frozen parameters approach described here https://docs.kidger.site/equinox/examples/frozen_layer/ , to setup up a training loop such that a different subset of parameters are frozen based upon the epoch

e.g For simplicity's sake, let's say you have an MLP, and you want the weights of odd and even layers to alternate between being updated each epoch (this isn't my real use case). As far as I can tell, if you use the filter_spec approach, the opt_state will have None values for the frozen parameters, so you will then run into issues when you apply the opposite filter. Am I missing something here?