src/modules/optim/sgd.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/sgd.h"

#include <vector>

#include "tenzr/common/error.h"

namespace tenzr::optim {

  SGD::SGD(double lr, double momentum, double dampening, double weight_decay, bool nesterov, ScheduleFn schedule)
      : lr_(lr), momentum_(momentum), dampening_(dampening), weight_decay_(weight_decay), nesterov_(nesterov), schedule_(std::move(schedule)) {
    TENZR_CHECK(lr >= 0.0, "Learning rate must be non-negative");
    TENZR_CHECK(momentum >= 0.0, "Momentum must be non-negative");
    TENZR_CHECK(weight_decay >= 0.0, "Weight decay must be non-negative");
  }

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

  std::unordered_map<std::string, Tensor> SGD::Update(const std::unordered_map<std::string, Tensor>& grads, OptState& state,
                                                      const std::unordered_map<std::string, Tensor>& params) {
    double current_lr = schedule_ ? schedule_(state.step) : lr_;
    double lr_d = current_lr;
    double mom_d = momentum_;
    double damp_d = dampening_;
    double wd_d = weight_decay_;

    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();
      std::vector<double> u_raw(n);

      if (momentum_ > 0.0) {
        auto& v_tensor = state.param_states[name]["v"];
        auto v_raw = v_tensor.ToVectorAs<double>();

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

          if (state.step == 0) {
            v_raw[i] = grad_val;
          } else {
            v_raw[i] = mom_d * v_raw[i] + (1.0 - damp_d) * grad_val;
          }

          double step_val = nesterov_ ? (grad_val + mom_d * v_raw[i]) : v_raw[i];
          u_raw[i] = -lr_d * step_val;
        }
        v_tensor.BindData(Tensor::FromVector<double>(v_raw, v_tensor.shape()).Astype(v_tensor.Dtype()).Data());
      } else {
        for (size_t i = 0; i < n; ++i) {
          double grad_val = g_raw[i];
          if (wd_d > 0.0) {
            grad_val += wd_d * p_raw[i];
          }
          u_raw[i] = -lr_d * grad_val;
        }
      }

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

    state.step++;
    return updates;
  }

}  // namespace tenzr::optim