optax.ScaleByAdaDeltaState