implement linear module

Signed-off-by: Ceng23333 <441651826@qq.com>

Author: Ceng23333

include/infinicore.hpp

@@ -1,4 +1,5 @@
 #pragma once
 
+#include "infinicore/nn.hpp"
 #include "infinicore/ops.hpp"
 #include "infinicore/tensor.hpp"

include/infinicore/nn.hpp

@@ -0,0 +1,3 @@
+#pragma once
+
+#include "nn/linear.hpp"

include/infinicore/nn/linear.hpp

@@ -0,0 +1,37 @@
+#pragma once
+
+#include "module.hpp"
+#include "../ops.hpp"
+
+namespace infinicore::nn {
+
+class Linear : public Module {
+public:
+    Linear(size_t in_features, size_t out_features, bool bias = true, const Device &device = Device());
+
+    // Forward pass: output = input @ weight.T + bias
+    Tensor forward(const Tensor &input) const;
+
+    // Forward pass with residual connection (InfiniLM-style)
+    // output = input @ weight.T + bias + residual
+    Tensor forward(const Tensor &input, const Tensor &residual) const;
+
+    // Accessors for parameters
+    Tensor weight() const;
+    Tensor bias() const;
+
+    // Module information
+    size_t in_features() const { return in_features_; }
+    size_t out_features() const { return out_features_; }
+    bool has_bias() const { return has_bias_; }
+
+    // String representation
+    std::string extra_repr() const;
+
+private:
+    size_t in_features_;
+    size_t out_features_;
+    bool has_bias_;
+};
+
+} // namespace infinicore::nn

include/infinicore/nn/module.hpp

@@ -17,6 +17,10 @@ class Module {
 
     Tensor register_parameter(const std::string &name, Parameter param);
 
+    // Create a Linear submodule-like parameter set (weight and optional bias)
+    // Mirrors torch.nn.Linear shapes: weight [out_features, in_features], bias [out_features]
+    void linear(const std::string &name, size_t in_features, size_t out_features, bool bias = true);
+
     template <typename M>
     std::shared_ptr<M> add_module(const std::string &name, std::shared_ptr<M> submodule) {
         submodules_[name] = submodule;

include/infinicore/tensor.hpp

@@ -146,6 +146,7 @@ class TensorImpl : public std::enable_shared_from_this<TensorImpl> {
 
     /**
      * Copy Data from another tensor to this tensor.
+     * Currently, only contigous tensors of the same dtype and shape are supported.
      *
      * @param src The source tensor to copy from
      *

src/infinicore-test/test_nn_module.cc

@@ -12,9 +12,9 @@
 namespace infinicore::test {
 
 // Simple test module that mimics torch.nn.Linear
-class TestLinearModule : public infinicore::nn::Module {
+class MockLinearModule : public infinicore::nn::Module {
 public:
-    TestLinearModule(int input_size, int output_size, const infinicore::Device &device)
+    MockLinearModule(int input_size, int output_size, const infinicore::Device &device)
         : input_size_(input_size), output_size_(output_size), device_(device) {
 
         // Initialize weight parameter (similar to torch.nn.Linear.weight)
@@ -71,6 +71,7 @@ class NNModuleTest : public MemoryTestFramework {
     TestResult testModuleHierarchy();
     TestResult testParameterLoading();
     TestResult testModuleComparison();
+    TestResult testModuleLinear();
 };
 
 } // namespace infinicore::test

src/infinicore-test/test_nn_module.h

@@ -0,0 +1,128 @@
+#include "infinicore/nn/linear.hpp"
+#include "infinicore/ops.hpp"
+#include <spdlog/spdlog.h>
+
+namespace infinicore::nn {
+
+Linear::Linear(size_t in_features, size_t out_features, bool bias, const Device &device)
+    : in_features_(in_features), out_features_(out_features), has_bias_(bias) {
+
+    device_ = device;
+
+    // Register weight parameter: [out_features, in_features]
+    register_parameter("weight", Parameter({out_features, in_features}, DataType::F32, device));
+
+    // Register bias parameter if requested: [out_features]
+    if (bias) {
+        register_parameter("bias", Parameter({out_features}, DataType::F32, device));
+    }
+
+    spdlog::debug("Created Linear module: in_features={}, out_features={}, bias={}",
+                  in_features, out_features, bias);
+}
+
+Tensor Linear::forward(const Tensor &input) const {
+    auto sd = state_dict();
+    auto weight = sd.at("weight");
+    auto bias_it = sd.find("bias");
+
+    // Create output tensor with shape [batch_size, out_features]
+    auto output_shape = input->shape();
+    output_shape[output_shape.size() - 1] = out_features_;
+    auto output = Tensor::empty(output_shape, input->dtype(), input->device());
+
+    // Transpose weight: [out_features, in_features] -> [in_features, out_features]
+    auto weight_t = weight->permute({1, 0});
+
+    // InfiniLM-style linear computation: output = input @ weight_t + bias
+    // Handle bias broadcasting similar to InferenceContext::linear
+    if (bias_it != sd.end()) {
+        auto bias = bias_it->second;
+
+        // Broadcast bias to output shape (similar to InfiniLM's bias handling)
+        size_t ndim_diff = output->ndim() - 1;
+        std::vector<Stride> strides(ndim_diff, 0);
+        strides.push_back(bias->stride(0));
+        auto bias_view = bias->as_strided(output->shape(), strides);
+
+        // First set output to bias (broadcasted)
+        infinicore::op::rearrange_(output, bias_view);
+
+        // Compute matmul result separately, then add to output
+        auto matmul_result = infinicore::op::matmul(input, weight_t);
+        infinicore::op::add_(output, output, matmul_result);
+    } else {
+        // No bias: just compute output = input @ weight_t
+        infinicore::op::matmul_(output, input, weight_t);
+    }
+
+    return output;
+}
+
+Tensor Linear::forward(const Tensor &input, const Tensor &residual) const {
+    auto sd = state_dict();
+    auto weight = sd.at("weight");
+    auto bias_it = sd.find("bias");
+
+    // Create output tensor with shape [batch_size, out_features]
+    auto output_shape = input->shape();
+    output_shape[output_shape.size() - 1] = out_features_;
+    auto output = Tensor::empty(output_shape, input->dtype(), input->device());
+
+    // Transpose weight: [out_features, in_features] -> [in_features, out_features]
+    auto weight_t = weight->permute({1, 0});
+
+    // InfiniLM-style computation with residual: output = input @ weight_t + bias + residual
+    if (bias_it != sd.end()) {
+        auto bias = bias_it->second;
+
+        // Broadcast bias to output shape
+        size_t ndim_diff = output->ndim() - 1;
+        std::vector<Stride> strides(ndim_diff, 0);
+        strides.push_back(bias->stride(0));
+        auto bias_view = bias->as_strided(output->shape(), strides);
+
+        // First set output to bias (broadcasted)
+        infinicore::op::rearrange_(output, bias_view);
+
+        // Compute matmul result separately, then add to output
+        auto matmul_result = infinicore::op::matmul(input, weight_t);
+        infinicore::op::add_(output, output, matmul_result);
+
+        // Add residual: output = output + residual
+        infinicore::op::add_(output, output, residual);
+    } else {
+        // No bias: compute output = input @ weight_t + residual
+        infinicore::op::matmul_(output, input, weight_t);
+        infinicore::op::add_(output, output, residual);
+    }
+
+    return output;
+}
+
+Tensor Linear::weight() const {
+    auto sd = state_dict();
+    auto it = sd.find("weight");
+    if (it != sd.end()) {
+        return it->second;
+    }
+    throw std::runtime_error("Weight parameter not found");
+}
+
+Tensor Linear::bias() const {
+    if (!has_bias_) {
+        throw std::runtime_error("Linear module does not have bias");
+    }
+    auto sd = state_dict();
+    auto it = sd.find("bias");
+    if (it != sd.end()) {
+        return it->second;
+    }
+    throw std::runtime_error("Bias parameter not found");
+}
+
+std::string Linear::extra_repr() const {
+    return "in_features=" + std::to_string(in_features_) + ", out_features=" + std::to_string(out_features_) + ", bias=" + (has_bias_ ? "true" : "false");
+}
+
+} // namespace infinicore::nn

src/infinicore/nn/linear.cc

@@ -51,3 +51,17 @@ void Module::collect_all_parameters(const std::string &prefix, std::unordered_ma
 }
 
 } // namespace infinicore::nn
+
+namespace infinicore::nn {
+void Module::linear(const std::string &name, size_t in_features, size_t out_features, bool bias) {
+    // Register weight parameter: [out_features, in_features]
+    register_parameter(name + ".weight",
+                       Parameter({out_features, in_features}, DataType::F32, device_));
+
+    // Register optional bias parameter: [out_features]
+    if (bias) {
+        register_parameter(name + ".bias",
+                           Parameter({out_features}, DataType::F32, device_));
+    }
+}
+} // namespace infinicore::nn