good first issue
Description
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.
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.