swiglu_module.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <optional>
#include "pulsatrix/linear_module.hpp"
#include "pulsatrix/module.hpp"
 
namespace pulsatrix {
 
class SwiGLUModule : public Module {
public:
    SwiGLUModule(int64_t d_model, int64_t d_ff, DeviceBackend* backend);
 
    [[nodiscard]] Tensor backward(const Tensor& grad_output) override;
 
    [[nodiscard]] OpType op_type() const override { return OpType::Elementwise; }
 
    [[nodiscard]] Tensor propagate_relevance(const Tensor& relevance_out, const LRPRuleConfig& config) override;
 
    [[nodiscard]] std::vector<NamedParamRef> named_parameters() override;
 
    [[nodiscard]] int64_t d_model() const { return d_model_; }
    [[nodiscard]] int64_t d_ff() const { return d_ff_; }
 
    [[nodiscard]] LinearModule& gate_proj() { return gate_proj_; }
    [[nodiscard]] LinearModule& up_proj() { return up_proj_; }
    [[nodiscard]] LinearModule& down_proj() { return down_proj_; }
 
 
    [[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 d_ff_;
    DeviceBackend* backend_;
 
    LinearModule gate_proj_;
    LinearModule up_proj_;
    LinearModule down_proj_;
 
    // Forward caches -- the operands the gate multiply's forward/backward/LRP rule needs.
    // gate_proj_/up_proj_/down_proj_ also cache their own inputs internally; this module
    // additionally caches these three for its own composed math, same choice
    // MultiHeadAttentionModule made for its Q/K/V/scores caches.
    Shape last_input_shape_ = Shape({0});
    int64_t last_n_flat_ = 0;
    Tensor last_gate_pre_;   
    Tensor last_gate_post_;  
    Tensor last_up_;         
    bool has_forwarded_ = false;
};
 
}  // namespace pulsatrix