retnet_module.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <optional>
#include <initializer_list>
#include <vector>
 
#include "pulsatrix/module.hpp"
 
namespace pulsatrix {
 
class RetNetModule : public Module {
public:
    RetNetModule(int64_t d_model, int64_t key_dim, float gamma, DeviceBackend* backend);
 
    [[nodiscard]] Tensor backward(const Tensor& grad_output) override;
 
    [[nodiscard]] OpType op_type() const override { return OpType::Recurrent; }
 
    void set_W_Q(std::initializer_list<float> values);
    void set_W_K(std::initializer_list<float> values);
    void set_W_V(std::initializer_list<float> values);
 
    void set_W_Q(const std::vector<float>& values);
    void set_W_K(const std::vector<float>& values);
    void set_W_V(const std::vector<float>& values);
 
    [[nodiscard]] const Tensor& W_Q() const { return w_q_; }
    [[nodiscard]] const Tensor& W_K() const { return w_k_; }
    [[nodiscard]] const Tensor& W_V() const { return w_v_; }
 
    [[nodiscard]] const Tensor& W_Q_grad() const { return w_q_grad_; }
    [[nodiscard]] const Tensor& W_K_grad() const { return w_k_grad_; }
    [[nodiscard]] const Tensor& W_V_grad() const { return w_v_grad_; }
 
    [[nodiscard]] float gamma() const { return gamma_; }
 
    [[nodiscard]] Tensor propagate_relevance(const Tensor& relevance_out, const LRPRuleConfig& config) override;
 
    [[nodiscard]] std::vector<NamedParamRef> named_parameters() override {
        return {{"w_q", {&w_q_, &w_q_grad_}}, {"w_k", {&w_k_, &w_k_grad_}}, {"w_v", {&w_v_, &w_v_grad_}}};
    }
 
 
    [[nodiscard]] std::optional<DeviceType> compute_device() const override { return backend_->device(); }
 
protected:
    [[nodiscard]] Tensor forward_impl(const Tensor& input) override;
 
private:
    int64_t d_model_;
    int64_t key_dim_;
    float gamma_;
    DeviceBackend* backend_;
    Tensor w_q_;  // (d_model, key_dim)
    Tensor w_k_;  // (d_model, key_dim)
    Tensor w_v_;  // (d_model, d_model)
    Tensor w_q_grad_;
    Tensor w_k_grad_;
    Tensor w_v_grad_;
    Tensor last_input_;   // (N, L, d_model)
    Tensor last_q_;       // (N, L, key_dim)
    Tensor last_k_;       // (N, L, key_dim)
    Tensor last_v_;       // (N, L, d_model)
    Tensor last_states_;  // (N, L+1, key_dim, d_model); index 0 along dim 1 = S_0 = 0
    int64_t last_L_ = 0;
    bool has_forwarded_ = false;
};
 
}  // namespace pulsatrix