Make FlaxLMSDiscreteScheduler jittable

[pcuenca]

Please, see discussion in #2155. I was under the wrong assumption that FlaxLMSDiscreteScheduler had been made jittable as part of #1661, but it never was.

[patrickvonplaten]

Thanks for opening an issue here!

[akshatvishu]

Hey! @patrickvonplaten , I will be willing to work on this as my issue but can you provide a rough summary of how to proceed or what the potential workflow looks like?

[dgarnitz]

@akshatvishu did this get fixed yet? If not, I am happy to jump into it.

@patrickvonplaten @pcuenca can this be fixed by adding @jax.jit to methods in FlaxLMSDiscreteScheduler? I imagine this can be tested by timing the jit & non-jit execution times to demonstrate jit is working - is that correct?

[Dan8991]

Since I don't see any open pull requests addressing this issue I guess that no-one is working on it anymore, can I give it a try?

[pcuenca]

@Dan8991 Absolutely!

[AbhinavJangra29]

is anyone working on it or is it fixed? if not can i start working on it?

[yiyixuxu]

@AbhinavJangra29 sure:)

[Dan8991]

Technically I still have the branch with the fix already implemented. It was never merged because the pull request remained opened for months without anyone reviewing it. If someone is willing to review it I can open a new one.