Patch optimizer rework

[benikm91]

I found three issues when integrating new optimizers into the GPT example...

1. LearningRateSchedule should be renamed to LearningRateScheduler

---

2. Adam constructor was buggy

class Adam[V: IsFloating](
    val learningRate: Tensor0[V],
    b1: Tensor0[V] = Tensor0(0.99), // This hard coded V as Float64, if not using Adam.of
    b2: Tensor0[V],
    epsilon: Tensor0[V] 

Solution: Remove default values and have a constructor with Double for the common case.

---

3. AdamState having V makes handling this state unnecessarily hard. I suggest hard-coding the beta to store and load as Float32.

case class TrainingState(
  params: GPT.Params[Float32],
  optState: LearningRateSchedulerState[GPT.Params[Float32], [P] =>> AdamState[P, Float32]], // <--
)

vs.

case class TrainingState(
  params: GPT.Params[Float32],
  optState: LearningRateSchedulerState[GPT.Params[Float32], AdamState], // <--
)

---

I hope this fixes it enough that we don't run into more issues.

[marcelluethi]

Looks good. Thanks for fixing it.