src/modules/nn/loss.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/nn/loss.h"

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

#include "tenzr/common/error.h"

namespace tenzr::nn {

  Tensor MseLoss(const Tensor& pred, const Tensor& target, std::string reduction) {
    TENZR_CHECK(pred.IsDefined() && target.IsDefined(), "Tensors must be defined");
    TENZR_CHECK(pred.shape() == target.shape(), "MSE loss shape mismatch: " + pred.shape().ToString() + " vs " + target.shape().ToString());

    Tensor p = pred.HasData() ? pred : pred.Eval();
    Tensor t = target.HasData() ? target : target.Eval();

    if (p.HasData() && t.HasData()) {
      auto p_vec = p.ToVectorAs<double>();
      auto t_vec = t.ToVectorAs<double>();
      size_t n = p_vec.size();

      std::vector<double> diff_sq(n);
      double sum = 0.0;
      for (size_t i = 0; i < n; ++i) {
        double d = p_vec[i] - t_vec[i];
        double sq = d * d;
        diff_sq[i] = sq;
        sum += sq;
      }

      if (reduction == "mean") {
        return Tensor::FromVector<double>({sum / static_cast<double>(n)}, Shape({1})).Astype(p.Dtype());
      } else if (reduction == "sum") {
        return Tensor::FromVector<double>({sum}, Shape({1})).Astype(p.Dtype());
      } else if (reduction == "none") {
        return Tensor::FromVector<double>(diff_sq, p.shape()).Astype(p.Dtype());
      } else {
        TENZR_THROW("Unsupported reduction: " + reduction);
      }
    }

    auto diff = pred - target;
    return diff * diff;
  }

  Tensor BceLoss(const Tensor& pred, const Tensor& target, double eps, std::string reduction) {
    TENZR_CHECK(pred.IsDefined() && target.IsDefined(), "Tensors must be defined");
    TENZR_CHECK(pred.shape() == target.shape(), "BCE loss shape mismatch: " + pred.shape().ToString() + " vs " + target.shape().ToString());

    Tensor p = pred.HasData() ? pred : pred.Eval();
    Tensor t = target.HasData() ? target : target.Eval();

    if (p.HasData() && t.HasData()) {
      auto p_vec = p.ToVectorAs<double>();
      auto t_vec = t.ToVectorAs<double>();
      size_t n = p_vec.size();

      std::vector<double> losses(n);
      double sum = 0.0;

      for (size_t i = 0; i < n; ++i) {
        double y = t_vec[i];
        double p_val = std::clamp(p_vec[i], eps, 1.0 - eps);
        double loss_val = -(y * std::log(p_val) + (1.0 - y) * std::log(1.0 - p_val));
        losses[i] = loss_val;
        sum += loss_val;
      }

      if (reduction == "mean") {
        return Tensor::FromVector<double>({sum / static_cast<double>(n)}, Shape({1})).Astype(pred.Dtype());
      } else if (reduction == "sum") {
        return Tensor::FromVector<double>({sum}, Shape({1})).Astype(pred.Dtype());
      } else if (reduction == "none") {
        return Tensor::FromVector<double>(losses, pred.shape()).Astype(pred.Dtype());
      } else {
        TENZR_THROW("Unsupported reduction: " + reduction);
      }
    }

    return Tensor::CustomOp("bce_loss", {pred, target});
  }

  Tensor cross_entropy_loss(const Tensor& logits, const Tensor& targets, double eps, std::string reduction) {
    (void)eps;
    TENZR_CHECK(logits.IsDefined() && targets.IsDefined(), "Tensors must be defined");
    TENZR_CHECK(logits.shape().Rank() == 2, "CrossEntropy logits must have rank 2 [B, C]");

    int64_t b = logits.shape().Dim(0);
    int64_t c = logits.shape().Dim(1);

    Tensor log = logits.HasData() ? logits : logits.Eval();
    Tensor tgt = targets.HasData() ? targets : targets.Eval();

    if (log.HasData() && tgt.HasData()) {
      auto logit_vec = log.ToVectorAs<double>();
      bool targets_are_indices = (tgt.shape().Rank() == 1 && tgt.shape().Dim(0) == b);
      bool targets_are_onehot = (tgt.shape().Rank() == 2 && tgt.shape().Dim(0) == b && tgt.shape().Dim(1) == c);

      TENZR_CHECK(targets_are_indices || targets_are_onehot, "Targets must be indices [B] or one-hot [B, C]");

      std::vector<double> losses(static_cast<size_t>(b));
      double sum = 0.0;

      if (targets_are_indices) {
        std::vector<int64_t> target_indices(static_cast<size_t>(b));
        if (targets.Dtype() == DType::kInt32) {
          auto raw = targets.ToVector<int32_t>();
          for (size_t i = 0; i < raw.size(); ++i) {
            target_indices[i] = raw[i];
          }
        } else if (targets.Dtype() == DType::kInt64) {
          target_indices = targets.ToVector<int64_t>();
        } else {
          target_indices = targets.ToVectorAs<int64_t>();
        }

        for (int64_t bi = 0; bi < b; ++bi) {
          size_t b_offset = static_cast<size_t>(bi * c);
          double max_val = logit_vec[b_offset];
          for (int64_t ci = 1; ci < c; ++ci) {
            max_val = std::max(max_val, logit_vec[b_offset + static_cast<size_t>(ci)]);
          }

          double sum_exp = 0.0;
          for (int64_t ci = 0; ci < c; ++ci) {
            sum_exp += std::exp(logit_vec[b_offset + static_cast<size_t>(ci)] - max_val);
          }
          double lse = max_val + std::log(sum_exp);

          int64_t target_c = target_indices[static_cast<size_t>(bi)];
          TENZR_CHECK(target_c >= 0 && target_c < c, "Target class index out of bounds");

          losses[static_cast<size_t>(bi)] = lse - logit_vec[b_offset + static_cast<size_t>(target_c)];
          sum += losses[static_cast<size_t>(bi)];
        }
      } else {
        auto target_vec = targets.ToVectorAs<double>();
        for (int64_t bi = 0; bi < b; ++bi) {
          size_t b_offset = static_cast<size_t>(bi * c);
          double max_val = logit_vec[b_offset];
          for (int64_t ci = 1; ci < c; ++ci) {
            max_val = std::max(max_val, logit_vec[b_offset + static_cast<size_t>(ci)]);
          }

          double sum_exp = 0.0;
          for (int64_t ci = 0; ci < c; ++ci) {
            sum_exp += std::exp(logit_vec[b_offset + static_cast<size_t>(ci)] - max_val);
          }
          double lse = max_val + std::log(sum_exp);

          double sample_loss = 0.0;
          for (int64_t ci = 0; ci < c; ++ci) {
            double p_true = target_vec[b_offset + static_cast<size_t>(ci)];
            sample_loss += p_true * (lse - logit_vec[b_offset + static_cast<size_t>(ci)]);
          }

          losses[static_cast<size_t>(bi)] = sample_loss;
          sum += sample_loss;
        }
      }

      if (reduction == "mean") {
        return Tensor::FromVector<double>({sum / static_cast<double>(b)}, Shape({1})).Astype(logits.Dtype());
      } else if (reduction == "sum") {
        return Tensor::FromVector<double>({sum}, Shape({1})).Astype(logits.Dtype());
      } else if (reduction == "none") {
        return Tensor::FromVector<double>(losses, Shape({b})).Astype(logits.Dtype());
      } else {
        TENZR_THROW("Unsupported reduction: " + reduction);
      }
    }

    return Tensor::CustomOp("cross_entropy_loss", {logits, targets});
  }

}  // namespace tenzr::nn