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

#include <cmath>
#include <vector>

#include "tenzr/common/error.h"
#include "tenzr/modules/random/random.h"

namespace tenzr::nn {

  Linear::Linear(int64_t in_features, int64_t out_features, bool use_bias, DType dtype, std::optional<uint64_t> seed)
      : in_features_(in_features), out_features_(out_features), use_bias_(use_bias) {
    TENZR_CHECK(in_features > 0, "in_features must be positive");
    TENZR_CHECK(out_features > 0, "out_features must be positive");

    if (seed.has_value()) {
      random::Generator gen(*seed);
      weight_ = RegisterParameter("weight", random::GlorotUniform({in_features_, out_features_}, in_features_, out_features_, dtype, gen));
    } else {
      weight_ =
          RegisterParameter("weight", random::GlorotUniform({in_features_, out_features_}, in_features_, out_features_, dtype, random::DefaultGenerator()));
    }

    if (use_bias_) {
      bias_ = RegisterParameter("bias", Tensor::Zeros({1, out_features_}, dtype));
    }
  }

  Tensor Linear::Forward(const Tensor& input) {
    TENZR_CHECK(input.IsDefined(), "Input tensor is uninitialized");
    auto out = input.Matmul(weight_);
    if (use_bias_) {
      if (out.shape().Rank() >= 2 && bias_.shape().Dim(0) == 1 && out.shape().Dim(0) > 1) {
        auto raw_bias = bias_.ToVectorAs<double>();
        int64_t batch = out.shape().Dim(0);
        std::vector<double> b_batch(static_cast<size_t>(batch * out_features_));
        for (int64_t b = 0; b < batch; ++b) {
          for (int64_t o = 0; o < out_features_; ++o) {
            b_batch[static_cast<size_t>(b * out_features_ + o)] = raw_bias[static_cast<size_t>(o)];
          }
        }
        bias_broadcast_ = Tensor::FromVector<double>(b_batch, out.shape()).Astype(bias_.Dtype());
        return out + bias_broadcast_;
      }
      return out + bias_;
    }
    return out;
  }

}  // namespace tenzr::nn