flatten_module.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <optional>
#include "pulsatrix/module.hpp"
 
namespace pulsatrix {
 
class FlattenModule : public Module {
public:
    explicit FlattenModule(DeviceBackend* backend) : backend_(backend), last_input_shape_(Shape({0})) {}
 
    [[nodiscard]] Tensor backward(const Tensor& grad_output) override {
        require_device(grad_output, backend_->device(), "FlattenModule::backward");
        Tensor grad_input(grad_output);
        grad_input.reshape(last_input_shape_);
        return grad_input;
    }
 
    [[nodiscard]] OpType op_type() const override { return OpType::Elementwise; }
 
    [[nodiscard]] Tensor propagate_relevance(const Tensor& relevance_out, const LRPRuleConfig&) override {
        require_device(relevance_out, backend_->device(), "FlattenModule::propagate_relevance");
        Tensor relevance_in(relevance_out);
        relevance_in.reshape(last_input_shape_);
        return relevance_in;
    }
 
    [[nodiscard]] bool supports_lrp_rule(LRPRule) const override { return true; }
 
 
    [[nodiscard]] std::optional<DeviceType> compute_device() const override { return backend_->device(); }
 
protected:
    [[nodiscard]] Tensor forward_impl(const Tensor& input) override {
        last_input_shape_ = input.shape();
        int64_t N = input.shape().dim(0);
        Tensor output(input);
        output.reshape(Shape({N, input.numel() / N}));
        return output;
    }
 
private:
    DeviceBackend* backend_;
    Shape last_input_shape_;
};
 
}  // namespace pulsatrix