mamba_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 MambaModule : public Module {
public:
    MambaModule(int64_t d_model, int64_t state_size, DeviceBackend* backend);
 
    [[nodiscard]] Tensor backward(const Tensor& grad_output) override;
 
    [[nodiscard]] OpType op_type() const override { return OpType::Recurrent; }
 
    void set_W_delta(std::initializer_list<float> values);
    void set_bias_delta(std::initializer_list<float> values);
    void set_W_B(std::initializer_list<float> values);
    void set_W_C(std::initializer_list<float> values);
    void set_A(std::initializer_list<float> values);
    void set_D(std::initializer_list<float> values);
 
    void set_W_delta(const std::vector<float>& values);
    void set_bias_delta(const std::vector<float>& values);
    void set_W_B(const std::vector<float>& values);
    void set_W_C(const std::vector<float>& values);
    void set_A(const std::vector<float>& values);
    void set_D(const std::vector<float>& values);
 
    [[nodiscard]] const Tensor& W_delta() const { return w_delta_; }
    [[nodiscard]] const Tensor& bias_delta() const { return bias_delta_; }
    [[nodiscard]] const Tensor& W_B() const { return w_b_; }
    [[nodiscard]] const Tensor& W_C() const { return w_c_; }
    [[nodiscard]] const Tensor& A() const { return a_; }
    [[nodiscard]] const Tensor& D() const { return d_; }
 
    [[nodiscard]] const Tensor& W_delta_grad() const { return w_delta_grad_; }
    [[nodiscard]] const Tensor& bias_delta_grad() const { return bias_delta_grad_; }
    [[nodiscard]] const Tensor& W_B_grad() const { return w_b_grad_; }
    [[nodiscard]] const Tensor& W_C_grad() const { return w_c_grad_; }
    [[nodiscard]] const Tensor& A_grad() const { return a_grad_; }
    [[nodiscard]] const Tensor& D_grad() const { return d_grad_; }
 
    [[nodiscard]] Tensor propagate_relevance(const Tensor& relevance_out, const LRPRuleConfig& config) override;
 
    [[nodiscard]] std::vector<NamedParamRef> named_parameters() override {
        return {
            {"w_delta", {&w_delta_, &w_delta_grad_}},
            {"bias_delta", {&bias_delta_, &bias_delta_grad_}},
            {"w_b", {&w_b_, &w_b_grad_}},
            {"w_c", {&w_c_, &w_c_grad_}},
            {"a", {&a_, &a_grad_}},
            {"d", {&d_, &d_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 state_size_;
    DeviceBackend* backend_;
    Tensor w_delta_;      // (d_model, d_model)
    Tensor bias_delta_;   // (d_model,)
    Tensor w_b_;          // (d_model, state_size)
    Tensor w_c_;          // (d_model, state_size)
    Tensor a_;            // (d_model, state_size)
    Tensor d_;            // (d_model,)
    Tensor w_delta_grad_;
    Tensor bias_delta_grad_;
    Tensor w_b_grad_;
    Tensor w_c_grad_;
    Tensor a_grad_;
    Tensor d_grad_;
    Tensor last_input_;    // (N, L, d_model)
    Tensor last_states_;   // (N, L+1, d_model, state_size); index 1 along dim 1 = h_0 = 0
    Tensor last_abar_;     // (N, L, d_model, state_size)
    Tensor last_bbar_;     // (N, L, d_model, state_size)
    Tensor last_z_delta_;  // (N, L, d_model); softplus PRE-activation, for its derivative
    Tensor last_delta_;    // (N, L, d_model)
    Tensor last_b_;        // (N, L, state_size)
    Tensor last_c_;        // (N, L, state_size)
    Tensor last_output_;   // (N, L, d_model); y_t, the LRP epsilon-rule denominator at step 1
    int64_t last_L_ = 0;
    bool has_forwarded_ = false;
};
 
}  // namespace pulsatrix