include/tenzr/modules/nn/norm.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_NN_NORM_H_
#define TENZR_MODULES_NN_NORM_H_

#include <cstdint>
#include <string>
#include <utility>
#include <vector>

#include "tenzr/frontend/dtype.h"
#include "tenzr/frontend/tensor.h"
#include "tenzr/modules/nn/module.h"

namespace tenzr::nn {

  class LayerNorm : public Module {
   public:
    // Dynamic DType constructor (explicit dtype)
    LayerNorm(int64_t normalized_dim, double eps, bool elementwise_affine, DType dtype);

    // Compile-time templated constructor using C++20 concepts (defaults to float via concepts::FloatingPoint<T>)
    template <typename T = float>
      requires concepts::FloatingPoint<T>
    explicit LayerNorm(int64_t normalized_dim, double eps = 1.0e-5, bool elementwise_affine = true)
        : LayerNorm(normalized_dim, eps, elementwise_affine, DtypeOf<T>()) {}

    // Tag dispatch constructor with explicit type: LayerNorm(std::type_identity<double>{}, 128)
    template <typename T>
      requires concepts::FloatingPoint<T>
    explicit LayerNorm(std::type_identity<T>, int64_t normalized_dim, double eps = 1.0e-5, bool elementwise_affine = true)
        : LayerNorm(normalized_dim, eps, elementwise_affine, DtypeOf<T>()) {}

    // Compile-time factory method: LayerNorm::create<double>(128)
    template <typename T = float>
      requires concepts::FloatingPoint<T>
    static LayerNorm Create(int64_t normalized_dim, double eps = 1.0e-5, bool elementwise_affine = true) {
      return LayerNorm(std::type_identity<T>{}, normalized_dim, eps, elementwise_affine);
    }

    Tensor Forward(const Tensor& input) override;

    Tensor& Weight() { return weight_; }
    const Tensor& Weight() const { return weight_; }

    Tensor& Bias() { return bias_; }
    const Tensor& Bias() const { return bias_; }

    int64_t NormalizedDim() const { return normalized_dim_; }
    double Eps() const { return eps_; }

    std::string Name() const override { return "LayerNorm"; }

   private:
    int64_t normalized_dim_;
    double eps_;
    bool elementwise_affine_;
    Tensor weight_;
    Tensor bias_;
  };

  class RMSNorm : public Module {
   public:
    // Dynamic DType constructor (explicit dtype)
    RMSNorm(int64_t dim, double eps, DType dtype);

    // Compile-time templated constructor using C++20 concepts (defaults to float via concepts::FloatingPoint<T>)
    template <typename T = float>
      requires concepts::FloatingPoint<T>
    explicit RMSNorm(int64_t dim, double eps = 1.0e-6) : RMSNorm(dim, eps, DtypeOf<T>()) {}

    // Tag dispatch constructor with explicit type: RMSNorm(std::type_identity<double>{}, 128)
    template <typename T>
      requires concepts::FloatingPoint<T>
    explicit RMSNorm(std::type_identity<T>, int64_t dim, double eps = 1.0e-6) : RMSNorm(dim, eps, DtypeOf<T>()) {}

    // Compile-time factory method: RMSNorm::create<double>(128)
    template <typename T = float>
      requires concepts::FloatingPoint<T>
    static RMSNorm Create(int64_t dim, double eps = 1.0e-6) {
      return RMSNorm(std::type_identity<T>{}, dim, eps);
    }

    Tensor Forward(const Tensor& input) override;

    Tensor& Weight() { return weight_; }
    const Tensor& Weight() const { return weight_; }

    int64_t Dim() const { return dim_; }
    double Eps() const { return eps_; }

    std::string Name() const override { return "RMSNorm"; }

   private:
    int64_t dim_;
    double eps_;
    Tensor weight_;
  };

  class BatchNorm : public Module {
   public:
    // Dynamic DType constructor (explicit dtype)
    BatchNorm(int64_t num_features, double eps, double momentum, DType dtype);

    // Compile-time templated constructor using C++20 concepts (defaults to float via concepts::FloatingPoint<T>)
    template <typename T = float>
      requires concepts::FloatingPoint<T>
    explicit BatchNorm(int64_t num_features, double eps = 1.0e-5, double momentum = 0.1) : BatchNorm(num_features, eps, momentum, DtypeOf<T>()) {}

    // Tag dispatch constructor with explicit type: BatchNorm(std::type_identity<double>{}, 128)
    template <typename T>
      requires concepts::FloatingPoint<T>
    explicit BatchNorm(std::type_identity<T>, int64_t num_features, double eps = 1.0e-5, double momentum = 0.1)
        : BatchNorm(num_features, eps, momentum, DtypeOf<T>()) {}

    // Compile-time factory method: BatchNorm::create<double>(128)
    template <typename T = float>
      requires concepts::FloatingPoint<T>
    static BatchNorm Create(int64_t num_features, double eps = 1.0e-5, double momentum = 0.1) {
      return BatchNorm(std::type_identity<T>{}, num_features, eps, momentum);
    }

    Tensor Forward(const Tensor& input) override;

    Tensor& Weight() { return weight_; }
    const Tensor& Weight() const { return weight_; }

    Tensor& Bias() { return bias_; }
    const Tensor& Bias() const { return bias_; }

    const std::vector<float>& RunningMean() const { return running_mean_; }
    const std::vector<float>& RunningVar() const { return running_var_; }

    int64_t NumFeatures() const { return num_features_; }
    double Eps() const { return eps_; }
    double Momentum() const { return momentum_; }

    std::string Name() const override { return "BatchNorm"; }

   private:
    int64_t num_features_;
    double eps_;
    double momentum_;
    Tensor weight_;
    Tensor bias_;
    std::vector<float> running_mean_;
    std::vector<float> running_var_;
  };

}  // namespace tenzr::nn

#endif  // TENZR_MODULES_NN_NORM_H_