make weight and bias fields of class && remove duplicate code

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

Author: Ceng23333

include/infinicore/nn/linear.hpp

@@ -17,8 +17,8 @@ class Linear : public Module {
     Tensor forward(const Tensor &input, const Tensor &residual) const;
 
     // Accessors for parameters
-    Tensor weight() const;
-    Tensor bias() const;
+    Tensor weight() const { return weight_; }
+    Tensor bias() const { return bias_; }
 
     // Module information
     size_t in_features() const { return in_features_; }
@@ -28,7 +28,14 @@ class Linear : public Module {
     // String representation
     std::string extra_repr() const;
 
+    // Direct access to parameters as fields
+    Parameter weight_;
+    Parameter bias_;
+
 private:
+    // Helper method for common forward computation
+    Tensor compute_linear(const Tensor &input) const;
+
     size_t in_features_;
     size_t out_features_;
     bool has_bias_;

src/infinicore/nn/linear.cc

@@ -5,45 +5,41 @@
 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) {
+    : weight_(Parameter({out_features, in_features}, DataType::F32, device)),
+      bias_(bias ? Parameter({out_features}, DataType::F32, device) : Parameter()),
+      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 weight parameter in state dict
+    register_parameter("weight", weight_);
 
-    // Register bias parameter if requested: [out_features]
+    // Register bias parameter if requested
     if (bias) {
-        register_parameter("bias", Parameter({out_features}, DataType::F32, device));
+        register_parameter("bias", bias_);
     }
 
     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");
-
+Tensor Linear::compute_linear(const Tensor &input) const {
     // 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;
+    auto weight_t = weight_->permute({1, 0});
 
-        // Broadcast bias to output shape (similar to InfiniLM's bias handling)
+    if (has_bias_) {
+        // 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);
+        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);
@@ -59,68 +55,19 @@ Tensor Linear::forward(const Tensor &input) const {
     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);
+Tensor Linear::forward(const Tensor &input) const {
+    return compute_linear(input);
+}
 
-        // Compute matmul result separately, then add to output
-        auto matmul_result = infinicore::op::matmul(input, weight_t);
-        infinicore::op::add_(output, output, matmul_result);
+Tensor Linear::forward(const Tensor &input, const Tensor &residual) const {
+    auto output = compute_linear(input);
 
-        // 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);
-    }
+    // Add residual: output = output + residual
+    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");
 }