include/tenzr/modules/nn/linear.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_LINEAR_H_
#define TENZR_MODULES_NN_LINEAR_H_

#include <cstdint>
#include <optional>
#include <string>

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

namespace tenzr::nn {

  class Linear : public Module {
   public:
    // Dynamic DType constructor (explicit dtype, optional seed)
    Linear(int64_t in_features, int64_t out_features, bool use_bias, DType dtype, std::optional<uint64_t> seed = std::nullopt);

    // Compile-time templated constructor using C++20 concepts (defaults to float via concepts::FloatingPoint<T>)
    template <typename T = float>
      requires concepts::FloatingPoint<T>
    Linear(int64_t in_features, int64_t out_features, bool use_bias = true, std::optional<uint64_t> seed = std::nullopt)
        : Linear(in_features, out_features, use_bias, DtypeOf<T>(), seed) {}

    // Tag dispatch constructor with explicit type: Linear(std::type_identity<double>{}, 8, 4)
    template <typename T>
      requires concepts::FloatingPoint<T>
    explicit Linear(std::type_identity<T>, int64_t in_features, int64_t out_features, bool use_bias = true, std::optional<uint64_t> seed = std::nullopt)
        : Linear(in_features, out_features, use_bias, DtypeOf<T>(), seed) {}

    // Compile-time factory method: Linear::create<double>(8, 4)
    template <typename T = float>
      requires concepts::FloatingPoint<T>
    static Linear Create(int64_t in_features, int64_t out_features, bool use_bias = true, std::optional<uint64_t> seed = std::nullopt) {
      return Linear(in_features, out_features, use_bias, DtypeOf<T>(), seed);
    }

    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 InFeatures() const { return in_features_; }
    int64_t OutFeatures() const { return out_features_; }
    bool HasBias() const { return use_bias_; }

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

   private:
    int64_t in_features_;
    int64_t out_features_;
    bool use_bias_;
    Tensor weight_;
    Tensor bias_;
    Tensor bias_broadcast_;
  };

}  // namespace tenzr::nn

#endif  // TENZR_MODULES_NN_LINEAR_H_