gru_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 GRUModule : public Module {
public:
    GRUModule(int64_t input_size, int64_t hidden_size, DeviceBackend* backend);
 
    [[nodiscard]] Tensor backward(const Tensor& grad_output) override;
 
    [[nodiscard]] OpType op_type() const override { return OpType::Recurrent; }
 
    void set_weight_xz(std::initializer_list<float> values);
    void set_weight_hz(std::initializer_list<float> values);
    void set_bias_z(std::initializer_list<float> values);
    void set_weight_xr(std::initializer_list<float> values);
    void set_weight_hr(std::initializer_list<float> values);
    void set_bias_r(std::initializer_list<float> values);
    void set_weight_xn(std::initializer_list<float> values);
    void set_weight_hn(std::initializer_list<float> values);
    void set_bias_n(std::initializer_list<float> values);
 
    [[nodiscard]] const Tensor& weight_xz() const { return weight_xz_; }
    [[nodiscard]] const Tensor& weight_hz() const { return weight_hz_; }
    [[nodiscard]] const Tensor& bias_z() const { return bias_z_; }
    [[nodiscard]] const Tensor& weight_xr() const { return weight_xr_; }
    [[nodiscard]] const Tensor& weight_hr() const { return weight_hr_; }
    [[nodiscard]] const Tensor& bias_r() const { return bias_r_; }
    [[nodiscard]] const Tensor& weight_xn() const { return weight_xn_; }
    [[nodiscard]] const Tensor& weight_hn() const { return weight_hn_; }
    [[nodiscard]] const Tensor& bias_n() const { return bias_n_; }
 
    [[nodiscard]] const Tensor& weight_xz_grad() const { return weight_xz_grad_; }
    [[nodiscard]] const Tensor& weight_hz_grad() const { return weight_hz_grad_; }
    [[nodiscard]] const Tensor& bias_z_grad() const { return bias_z_grad_; }
    [[nodiscard]] const Tensor& weight_xr_grad() const { return weight_xr_grad_; }
    [[nodiscard]] const Tensor& weight_hr_grad() const { return weight_hr_grad_; }
    [[nodiscard]] const Tensor& bias_r_grad() const { return bias_r_grad_; }
    [[nodiscard]] const Tensor& weight_xn_grad() const { return weight_xn_grad_; }
    [[nodiscard]] const Tensor& weight_hn_grad() const { return weight_hn_grad_; }
    [[nodiscard]] const Tensor& bias_n_grad() const { return bias_n_grad_; }
 
    [[nodiscard]] Tensor propagate_relevance(const Tensor& relevance_out, const LRPRuleConfig& config) override;
 
    [[nodiscard]] std::vector<NamedParamRef> named_parameters() override {
        return {
            {"weight_xz", {&weight_xz_, &weight_xz_grad_}},
            {"weight_hz", {&weight_hz_, &weight_hz_grad_}},
            {"bias_z", {&bias_z_, &bias_z_grad_}},
            {"weight_xr", {&weight_xr_, &weight_xr_grad_}},
            {"weight_hr", {&weight_hr_, &weight_hr_grad_}},
            {"bias_r", {&bias_r_, &bias_r_grad_}},
            {"weight_xn", {&weight_xn_, &weight_xn_grad_}},
            {"weight_hn", {&weight_hn_, &weight_hn_grad_}},
            {"bias_n", {&bias_n_, &bias_n_grad_}},
        };
    }
 
 
    [[nodiscard]] std::optional<DeviceType> compute_device() const override { return backend_->device(); }
 
protected:
    [[nodiscard]] Tensor forward_impl(const Tensor& input) override;
 
private:
    int64_t input_size_;
    int64_t hidden_size_;
    DeviceBackend* backend_;
    Tensor weight_xz_;  // (input_size, hidden_size)
    Tensor weight_hz_;  // (hidden_size, hidden_size)
    Tensor bias_z_;     // (hidden_size,)
    Tensor weight_xr_;
    Tensor weight_hr_;
    Tensor bias_r_;
    Tensor weight_xn_;
    Tensor weight_hn_;
    Tensor bias_n_;
    Tensor weight_xz_grad_;
    Tensor weight_hz_grad_;
    Tensor bias_z_grad_;
    Tensor weight_xr_grad_;
    Tensor weight_hr_grad_;
    Tensor bias_r_grad_;
    Tensor weight_xn_grad_;
    Tensor weight_hn_grad_;
    Tensor bias_n_grad_;
    Tensor last_input_;           // (N, L, input_size)
    Tensor last_hidden_states_;   // (N, L+1, hidden_size); index 0 = h_0 = 0
    Tensor last_gate_z_;          // (N, L, hidden_size); update gate
    Tensor last_gate_r_;          // (N, L, hidden_size); reset gate
    Tensor last_candidate_n_;     // (N, L, hidden_size); n_t, the tanh candidate signal
    // hn_prev_t = h_{t-1} @ W_hn is a DISTINCT intermediate from h_{t-1} itself (the reset
    // gate multiplies the projection, not the raw state), and it is the node the candidate
    // path's relevance lands on before being redistributed back onto h_{t-1} -- so it is
    // cached in its own right, for both backward() and propagate_relevance()'s step 6.
    Tensor last_hn_prev_;         // (N, L, hidden_size)
    Tensor last_pre_activation_n_;  // (N, L, hidden_size); n_t's pre-activation EXCLUDING bias
    int64_t last_L_ = 0;
    bool has_forwarded_ = false;
};
 
}  // namespace pulsatrix