spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2#pragma once34#include "core/autograd.h"5#include "ops/cpu/cpu_ops.h"67#include <random>8#include <string>9#include <utility>10#include <vector>1112namespace forge::nn {1314// Base module: owns nothing but the parameter registry. Children register15// their parameters at construction; parents absorb them with a name prefix.16// Tied parameters may appear under several names — consumers that update17// parameters (optimizer, checkpoint) must dedupe by Var::id().18class Module {19public:20    virtual ~Module() = default;2122    const std::vector<std::pair<std::string, Var>>& named_parameters() const {23        return params_;24    }2526    void zero_grad() const {27        for (const auto& [name, p] : params_) p.zero_grad();28    }2930protected:31    Var register_param(const std::string& name, Tensor init) {32        Var p(std::move(init), /*requires_grad=*/true);33        params_.emplace_back(name, p);34        return p;35    }3637    // Register an existing Var under a new name (weight tying).38    Var register_param(const std::string& name, const Var& shared) {39        params_.emplace_back(name, shared);40        return shared;41    }4243    void absorb(const std::string& prefix, const Module& child) {44        for (const auto& [name, p] : child.named_parameters())45            params_.emplace_back(prefix + "." + name, p);46    }4748    std::vector<std::pair<std::string, Var>> params_;49};5051// init helper shared by all modules52inline Tensor normal_init(std::vector<int64_t> shape, float stddev, std::mt19937_64& rng) {53    Tensor t = Tensor::empty(std::move(shape));54    cpu::fill_normal(t, 0.0f, stddev, rng);55    return t;56}5758} // namespace forge::nn59