embedding_module.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <initializer_list>
#include <vector>
 
#include "pulsatrix/module.hpp"
 
namespace pulsatrix {
 
class EmbeddingModule : public Module {
public:
    EmbeddingModule(int64_t num_embeddings, int64_t embedding_dim, DeviceBackend* backend);
 
    [[nodiscard]] Tensor backward(const Tensor& grad_output) override;
 
    [[nodiscard]] OpType op_type() const override { return OpType::Embedding; }
 
    void set_weight(std::initializer_list<float> values);
    void set_weight(const std::vector<float>& values);
 
    [[nodiscard]] const Tensor& weight() const { return weight_; }
    [[nodiscard]] const Tensor& weight_grad() const { return weight_grad_; }
 
    [[nodiscard]] Tensor propagate_relevance(const Tensor& relevance_out, const LRPRuleConfig& config) override;
 
    [[nodiscard]] std::vector<NamedParamRef> named_parameters() override {
        return {
            {"weight", {&weight_, &weight_grad_}},
        };
    }
 
protected:
    [[nodiscard]] Tensor forward_impl(const Tensor& input) override;
 
private:
    int64_t num_embeddings_;
    int64_t embedding_dim_;
    DeviceBackend* backend_;
    Tensor weight_;       // shape (num_embeddings, embedding_dim)
    Tensor weight_grad_;
    Shape last_input_shape_ = Shape({0});
    // Flat N*L validated indices as whole-number floats on the weights' device (exact below
    // 2^24 -- far above any vocabulary this module is built for), consumed by gather_rows /
    // scatter_add_rows (GPU-native-kernels Mission 2).
    Tensor last_indices_ = Tensor(Shape({0}), backend_);
    bool has_forwarded_ = false;
};
 
}  // namespace pulsatrix