src/modules/optim/adam.cc¶

// 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.

#include "tenzr/modules/optim/adam.h"

#include <algorithm>
#include <cmath>
#include <vector>

#include "tenzr/common/error.h"

namespace tenzr::optim {

  Adam::Adam(double lr, double beta1, double beta2, double eps, double weight_decay, bool amsgrad, bool decoupled_weight_decay, ScheduleFn schedule)
      : lr_(lr),
        beta1_(beta1),
        beta2_(beta2),
        eps_(eps),
        weight_decay_(weight_decay),
        amsgrad_(amsgrad),
        decoupled_weight_decay_(decoupled_weight_decay),
        schedule_(std::move(schedule)) {
    TENZR_CHECK(lr >= 0.0, "Learning rate must be non-negative");
    TENZR_CHECK(beta1 >= 0.0 && beta1 < 1.0, "Beta1 must be in [0, 1)");
    TENZR_CHECK(beta2 >= 0.0 && beta2 < 1.0, "Beta2 must be in [0, 1)");
    TENZR_CHECK(eps > 0.0, "Eps must be positive");
  }

  OptState Adam::Init(const std::unordered_map<std::string, Tensor>& params) {
    OptState state;
    state.step = 0;
    for (const auto& [name, p] : params) {
      state.param_states[name]["m"] = Tensor::Zeros(p.shape(), p.Dtype());
      state.param_states[name]["v"] = Tensor::Zeros(p.shape(), p.Dtype());
      if (amsgrad_) {
        state.param_states[name]["v_hat_max"] = Tensor::Zeros(p.shape(), p.Dtype());
      }
    }
    return state;
  }

  std::unordered_map<std::string, Tensor> Adam::Update(const std::unordered_map<std::string, Tensor>& grads, OptState& state,
                                                       const std::unordered_map<std::string, Tensor>& params) {
    uint64_t step = state.step + 1;
    double current_lr = schedule_ ? schedule_(state.step) : lr_;

    double bc1 = 1.0 - std::pow(beta1_, static_cast<double>(step));
    double bc2 = 1.0 - std::pow(beta2_, static_cast<double>(step));

    double lr_d = current_lr;
    double b1_d = beta1_;
    double b2_d = beta2_;
    double eps_d = eps_;
    double wd_d = weight_decay_;
    double bc1_d = bc1;
    double bc2_sqrt_d = std::sqrt(bc2);

    std::unordered_map<std::string, Tensor> updates;

    for (const auto& [name, g] : grads) {
      auto p_it = params.find(name);
      TENZR_CHECK(p_it != params.end(), "Gradient key not found in parameters: " + name);
      const auto& p = p_it->second;

      // A lazy gradient is consumed here (to_vector materializes it).
      if (!p.HasData()) {
        updates[name] = g;
        continue;
      }

      auto g_raw = g.ToVectorAs<double>();
      auto p_raw = p.ToVectorAs<double>();
      size_t n = g_raw.size();

      auto& m_tensor = state.param_states[name]["m"];
      auto& v_tensor = state.param_states[name]["v"];
      auto m_raw = m_tensor.ToVectorAs<double>();
      auto v_raw = v_tensor.ToVectorAs<double>();

      std::vector<double> v_hat_max_raw;
      if (amsgrad_) {
        v_hat_max_raw = state.param_states[name]["v_hat_max"].ToVectorAs<double>();
      }

      std::vector<double> u_raw(n);

      for (size_t i = 0; i < n; ++i) {
        double grad_val = g_raw[i];
        if (!decoupled_weight_decay_ && wd_d > 0.0) {
          grad_val += wd_d * p_raw[i];
        }

        // Update moments
        m_raw[i] = b1_d * m_raw[i] + (1.0 - b1_d) * grad_val;
        v_raw[i] = b2_d * v_raw[i] + (1.0 - b2_d) * (grad_val * grad_val);

        double v_denom;
        if (amsgrad_) {
          v_hat_max_raw[i] = std::max(v_hat_max_raw[i], v_raw[i]);
          v_denom = (std::sqrt(v_hat_max_raw[i]) / bc2_sqrt_d) + eps_d;
        } else {
          v_denom = (std::sqrt(v_raw[i]) / bc2_sqrt_d) + eps_d;
        }

        double m_hat = m_raw[i] / bc1_d;
        double step_val = -lr_d * (m_hat / v_denom);

        if (decoupled_weight_decay_ && wd_d > 0.0) {
          step_val -= lr_d * wd_d * p_raw[i];
        }

        u_raw[i] = step_val;
      }

      m_tensor.BindData(Tensor::FromVector<double>(m_raw, m_tensor.shape()).Astype(m_tensor.Dtype()).Data());
      v_tensor.BindData(Tensor::FromVector<double>(v_raw, v_tensor.shape()).Astype(v_tensor.Dtype()).Data());
      if (amsgrad_) {
        auto& vmax = state.param_states[name]["v_hat_max"];
        vmax.BindData(Tensor::FromVector<double>(v_hat_max_raw, vmax.shape()).Astype(vmax.Dtype()).Data());
      }

      updates[name] = Tensor::FromVector<double>(u_raw, p.shape()).Astype(p.Dtype());
    }

    state.step = step;
    return updates;
  }

}  // namespace tenzr::optim