// 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/norm.h"
#include <cmath>
#include <vector>
#include "tenzr/common/error.h"
namespace tenzr::nn {
// LayerNorm
LayerNorm::LayerNorm(int64_t normalized_dim, double eps, bool elementwise_affine, DType dtype)
: normalized_dim_(normalized_dim), eps_(eps), elementwise_affine_(elementwise_affine) {
TENZR_CHECK(normalized_dim > 0, "normalized_dim must be positive");
if (elementwise_affine_) {
auto w = Tensor::Ones({normalized_dim_}, dtype);
weight_ = RegisterParameter("weight", w);
auto b = Tensor::Zeros({normalized_dim_}, dtype);
bias_ = RegisterParameter("bias", b);
}
}
Tensor LayerNorm::Forward(const Tensor& input) {
TENZR_CHECK(input.IsDefined(), "Input tensor is uninitialized");
int64_t last_dim = input.shape().Dim(input.shape().Rank() - 1);
TENZR_CHECK(last_dim == normalized_dim_,
"Last dimension of input (" + std::to_string(last_dim) + ") must match normalized_dim (" + std::to_string(normalized_dim_) + ")");
if (input.HasData()) {
auto in_vec = input.ToVectorAs<double>();
int64_t total_elems = input.shape().NumElements();
int64_t d = normalized_dim_;
int64_t num_rows = total_elems / d;
std::vector<double> w_vec;
std::vector<double> b_vec;
if (elementwise_affine_ && weight_.HasData() && bias_.HasData()) {
w_vec = weight_.ToVectorAs<double>();
b_vec = bias_.ToVectorAs<double>();
}
std::vector<double> out_vec(static_cast<size_t>(total_elems));
for (int64_t r = 0; r < num_rows; ++r) {
size_t offset = static_cast<size_t>(r * d);
double mean = 0.0;
for (int64_t j = 0; j < d; ++j) {
mean += in_vec[offset + static_cast<size_t>(j)];
}
mean /= static_cast<double>(d);
double var = 0.0;
for (int64_t j = 0; j < d; ++j) {
double diff = in_vec[offset + static_cast<size_t>(j)] - mean;
var += diff * diff;
}
var /= static_cast<double>(d);
double inv_std = 1.0 / std::sqrt(var + eps_);
for (int64_t j = 0; j < d; ++j) {
double norm_val = (in_vec[offset + static_cast<size_t>(j)] - mean) * inv_std;
if (elementwise_affine_ && !w_vec.empty() && !b_vec.empty()) {
norm_val = norm_val * w_vec[static_cast<size_t>(j)] + b_vec[static_cast<size_t>(j)];
}
out_vec[offset + static_cast<size_t>(j)] = norm_val;
}
}
return Tensor::FromVector<double>(out_vec, input.shape()).Astype(input.Dtype());
}
auto mean = input.Mean(std::vector<int64_t>{-1}, /*keepdims=*/true);
auto diff = input - mean;
auto var = (diff * diff).Mean(std::vector<int64_t>{-1}, /*keepdims=*/true);
auto eps_arr = Array::Full(var.shape(), eps_, var.Dtype());
auto inv_std = (var + eps_arr).Rsqrt();
auto norm = diff * inv_std;
if (elementwise_affine_) {
norm = norm * weight_ + bias_;
}
return norm;
}
// RMSNorm
RMSNorm::RMSNorm(int64_t dim, double eps, DType dtype) : dim_(dim), eps_(eps) {
TENZR_CHECK(dim > 0, "dim must be positive");
auto w = Tensor::Ones({dim_}, dtype);
weight_ = RegisterParameter("weight", w);
}
Tensor RMSNorm::Forward(const Tensor& input) {
TENZR_CHECK(input.IsDefined(), "Input tensor is uninitialized");
int64_t last_dim = input.shape().Dim(input.shape().Rank() - 1);
TENZR_CHECK(last_dim == dim_, "Last dimension of input (" + std::to_string(last_dim) + ") must match RMSNorm dim (" + std::to_string(dim_) + ")");
if (input.HasData() && weight_.HasData()) {
auto in_vec = input.ToVectorAs<double>();
auto w_vec = weight_.ToVectorAs<double>();
int64_t total_elems = input.shape().NumElements();
int64_t d = dim_;
int64_t num_rows = total_elems / d;
std::vector<double> out_vec(static_cast<size_t>(total_elems));
for (int64_t r = 0; r < num_rows; ++r) {
size_t offset = static_cast<size_t>(r * d);
double sq_sum = 0.0;
for (int64_t j = 0; j < d; ++j) {
double val = in_vec[offset + static_cast<size_t>(j)];
sq_sum += val * val;
}
double rms = std::sqrt(sq_sum / static_cast<double>(d) + eps_);
double inv_rms = 1.0 / rms;
for (int64_t j = 0; j < d; ++j) {
out_vec[offset + static_cast<size_t>(j)] = in_vec[offset + static_cast<size_t>(j)] * inv_rms * w_vec[static_cast<size_t>(j)];
}
}
return Tensor::FromVector<double>(out_vec, input.shape()).Astype(input.Dtype());
}
auto sq_mean = (input * input).Mean(std::vector<int64_t>{-1}, /*keepdims=*/true);
auto eps_arr = Array::Full(sq_mean.shape(), eps_, sq_mean.Dtype());
auto inv_rms = (sq_mean + eps_arr).Rsqrt();
return input * inv_rms * weight_;
}
// BatchNorm
BatchNorm::BatchNorm(int64_t num_features, double eps, double momentum, DType dtype)
: num_features_(num_features),
eps_(eps),
momentum_(momentum),
running_mean_(static_cast<size_t>(num_features), 0.0f),
running_var_(static_cast<size_t>(num_features), 1.0f) {
TENZR_CHECK(num_features > 0, "num_features must be positive");
auto w = Tensor::Ones({num_features_}, dtype);
weight_ = RegisterParameter("weight", w);
auto b = Tensor::Zeros({num_features_}, dtype);
bias_ = RegisterParameter("bias", b);
}
Tensor BatchNorm::Forward(const Tensor& input) {
TENZR_CHECK(input.IsDefined(), "Input tensor is uninitialized");
TENZR_CHECK(input.shape().Rank() >= 2, "BatchNorm requires rank >= 2");
TENZR_CHECK(input.shape().Dim(1) == num_features_,
"Channel dimension (" + std::to_string(input.shape().Dim(1)) + ") must match num_features (" + std::to_string(num_features_) + ")");
if (input.HasData() && weight_.HasData() && bias_.HasData()) {
auto in_vec = input.ToVectorAs<double>();
auto w_vec = weight_.ToVectorAs<double>();
auto b_vec = bias_.ToVectorAs<double>();
int64_t b = input.shape().Dim(0);
int64_t c = num_features_;
int64_t spatial = input.shape().NumElements() / (b * c);
std::vector<double> out_vec(static_cast<size_t>(input.shape().NumElements()));
if (IsTraining()) {
// Compute batch mean and var per channel
int64_t count = b * spatial;
for (int64_t ci = 0; ci < c; ++ci) {
double ch_mean = 0.0;
for (int64_t bi = 0; bi < b; ++bi) {
for (int64_t s = 0; s < spatial; ++s) {
size_t idx = static_cast<size_t>((bi * c + ci) * spatial + s);
ch_mean += in_vec[idx];
}
}
ch_mean /= static_cast<double>(count);
double ch_var = 0.0;
for (int64_t bi = 0; bi < b; ++bi) {
for (int64_t s = 0; s < spatial; ++s) {
size_t idx = static_cast<size_t>((bi * c + ci) * spatial + s);
double diff = in_vec[idx] - ch_mean;
ch_var += diff * diff;
}
}
ch_var /= static_cast<double>(count);
// Update running statistics
running_mean_[static_cast<size_t>(ci)] =
static_cast<float>((1.0 - momentum_) * static_cast<double>(running_mean_[static_cast<size_t>(ci)]) + momentum_ * ch_mean);
running_var_[static_cast<size_t>(ci)] =
static_cast<float>((1.0 - momentum_) * static_cast<double>(running_var_[static_cast<size_t>(ci)]) + momentum_ * ch_var);
double inv_std = 1.0 / std::sqrt(ch_var + eps_);
double gamma = w_vec[static_cast<size_t>(ci)];
double beta = b_vec[static_cast<size_t>(ci)];
for (int64_t bi = 0; bi < b; ++bi) {
for (int64_t s = 0; s < spatial; ++s) {
size_t idx = static_cast<size_t>((bi * c + ci) * spatial + s);
out_vec[idx] = (in_vec[idx] - ch_mean) * inv_std * gamma + beta;
}
}
}
} else {
// Evaluation mode using running stats
for (int64_t ci = 0; ci < c; ++ci) {
double ch_mean = running_mean_[static_cast<size_t>(ci)];
double ch_var = running_var_[static_cast<size_t>(ci)];
double inv_std = 1.0 / std::sqrt(ch_var + eps_);
double gamma = w_vec[static_cast<size_t>(ci)];
double beta = b_vec[static_cast<size_t>(ci)];
for (int64_t bi = 0; bi < b; ++bi) {
for (int64_t s = 0; s < spatial; ++s) {
size_t idx = static_cast<size_t>((bi * c + ci) * spatial + s);
out_vec[idx] = (in_vec[idx] - ch_mean) * inv_std * gamma + beta;
}
}
}
}
return Tensor::FromVector<double>(out_vec, input.shape()).Astype(input.Dtype());
}
return Tensor::CustomOp("batchnorm", {input, weight_, bias_});
}
} // namespace tenzr::nn