spb/forge

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2#pragma once34#include "nn/module.h"5#include "ops/ops.h"67namespace forge::nn {89// Token (or positional) embedding table [V, C].10class Embedding : public Module {11public:12    Embedding(int64_t num_embeddings, int64_t dim, float weight_std, std::mt19937_64& rng) {13        weight_ = register_param("weight", normal_init({num_embeddings, dim}, weight_std, rng));14    }1516    // ids: [B, T] u16/i32 → [B, T, C]17    Var forward(const Tensor& ids) const { return ops::embedding(weight_, ids); }1819    const Var& weight() const { return weight_; }2021private:22    Var weight_;23};2425} // namespace forge::nn26