Different learning rates for different parts of an array

[ajithmoola]

I have a parameter of size (1000) and I want to use different learning rates for different parts of the array. For example, 0:200 I want to assign a learning rate of 0.1, from 200:600 learning rate should be 0.01 etc. Is this implementable in optax?

[fabianp]

Hi @ajithmoola . I don't believe we have such a transformation yet, although it seems to me that this could be useful to have.

What do you think @vroulet @mtthss ?

[vroulet]

I think https://optax.readthedocs.io/en/latest/api/combining_optimizers.html#multi-transform is what you are searching for