// Author: Simon-Pierre Boucher — contact@spboucher.ai
#pragma once

#include "nn/module.h"
#include "ops/ops.h"

namespace forge::nn {

// Token (or positional) embedding table [V, C].
class Embedding : public Module {
public:
    Embedding(int64_t num_embeddings, int64_t dim, float weight_std, std::mt19937_64& rng) {
        weight_ = register_param("weight", normal_init({num_embeddings, dim}, weight_std, rng));
    }

    // ids: [B, T] u16/i32 → [B, T, C]
    Var forward(const Tensor& ids) const { return ops::embedding(weight_, ids); }

    const Var& weight() const { return weight_; }

private:
    Var weight_;
};

} // namespace forge::nn