spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2#pragma once34#include "nn/module.h"5#include "ops/ops.h"67#include <string>89namespace forge::nn {1011inline ops::QuantMode quant_mode_from(const std::string& s) {12    if (s == "int8") return ops::QuantMode::Int8;13    if (s == "ternary") return ops::QuantMode::Ternary;14    return ops::QuantMode::None;15}1617// y = x ⋅ Wᵀ (+ b). Weight stored [out, in] (PyTorch convention), x is18// [N, in] 2-D — callers flatten [B,T,C] with Var::reshaped.19//20// Research seams (see CLAUDE.md):21//  - set_mask(): optional binary mask multiplied into the weight each22//    forward (lottery-ticket / pruning). Gradients of masked weights are23//    zeroed by the same multiply.24//  - forward() is virtual so quantized-weight variants (BitNet-style) can25//    override without touching call sites.26class Linear : public Module {27public:28    Linear(int64_t in_features, int64_t out_features, bool has_bias, float weight_std,29           std::mt19937_64& rng) {30        weight_ = register_param("weight",31                                 normal_init({out_features, in_features}, weight_std, rng));32        if (has_bias) bias_ = register_param("bias", Tensor::zeros({out_features}));33    }3435    virtual ~Linear() = default;3637    virtual Var forward(const Var& x) const;3839    // mask: [out, in], 0/1 f32. Undefined tensor clears the mask.40    void set_mask(Tensor mask) {41        mask_ = mask.defined() ? Var(std::move(mask), /*requires_grad=*/false) : Var();42    }4344    // QAT: fake-quantize the weight each forward (STE backward). The master45    // weight and optimizer state stay f32.46    void set_quant(ops::QuantMode mode) { quant_ = mode; }4748    const Var& weight() const { return weight_; }4950protected:51    Var weight_;52    Var bias_; // undefined if bias disabled53    Var mask_; // undefined if no mask54    ops::QuantMode quant_ = ops::QuantMode::None;55};5657} // namespace forge::nn58