include/tenzr/modules/optim/schedule.h¶

// Copyright 2026 The Tenzr Project Authors
//
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef TENZR_MODULES_OPTIM_SCHEDULE_H_
#define TENZR_MODULES_OPTIM_SCHEDULE_H_

#include <algorithm>
#include <cmath>
#include <cstdint>
#include <functional>

namespace tenzr::optim {

  using ScheduleFn = std::function<double(uint64_t step)>;

  inline ScheduleFn ConstantSchedule(double lr) {
    return [lr](uint64_t /*step*/) -> double { return lr; };
  }

  inline ScheduleFn LinearWarmup(double peak_lr, uint64_t warmup_steps) {
    return [peak_lr, warmup_steps](uint64_t step) -> double {
      if (warmup_steps == 0) {
        return peak_lr;
      }
      if (step >= warmup_steps) {
        return peak_lr;
      }
      return peak_lr * (static_cast<double>(step) / static_cast<double>(warmup_steps));
    };
  }

  inline ScheduleFn CosineDecay(double init_lr, uint64_t decay_steps, double alpha = 0.0) {
    constexpr double kPi = 3.14159265358979323846;
    return [init_lr, decay_steps, alpha](uint64_t step) -> double {
      if (step >= decay_steps) {
        return init_lr * alpha;
      }
      double progress = static_cast<double>(step) / static_cast<double>(decay_steps);
      double cosine = 0.5 * (1.0 + std::cos(kPi * progress));
      return init_lr * ((1.0 - alpha) * cosine + alpha);
    };
  }

  inline ScheduleFn WarmupCosineDecay(double peak_lr, uint64_t warmup_steps, uint64_t total_steps, double min_lr = 0.0) {
    constexpr double kPi = 3.14159265358979323846;
    return [peak_lr, warmup_steps, total_steps, min_lr](uint64_t step) -> double {
      if (step < warmup_steps) {
        return peak_lr * (static_cast<double>(step) / static_cast<double>(warmup_steps));
      }
      if (step >= total_steps) {
        return min_lr;
      }
      uint64_t decay_steps = total_steps - warmup_steps;
      uint64_t current_decay = step - warmup_steps;
      double progress = static_cast<double>(current_decay) / static_cast<double>(decay_steps);
      double cosine = 0.5 * (1.0 + std::cos(kPi * progress));
      return min_lr + (peak_lr - min_lr) * cosine;
    };
  }

}  // namespace tenzr::optim

#endif  // TENZR_MODULES_OPTIM_SCHEDULE_H_