activation_snapshot.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <optional>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
 
#include "pulsatrix/assert.hpp"
#include "pulsatrix/node.hpp"
#include "pulsatrix/op_type.hpp"
#include "pulsatrix/tensor.hpp"
 
namespace pulsatrix {
 
class ActivationSnapshot {
public:
    struct NodeMetadata {
        OpType op_type;
 
        std::optional<std::string> label;
    };
 
    ActivationSnapshot(std::unordered_map<NodeId, Tensor> activations, std::vector<NodeId> node_ids,
                       std::unordered_map<NodeId, NodeMetadata> metadata)
        : activations_(std::move(activations)), node_ids_(std::move(node_ids)), metadata_(std::move(metadata)) {}
 
    [[nodiscard]] const Tensor& activation(NodeId id) const {
        auto it = activations_.find(id);
        PULSATRIX_ASSERT(it != activations_.end());
        return it->second;
    }
 
    [[nodiscard]] const std::vector<NodeId>& node_ids() const { return node_ids_; }
 
    [[nodiscard]] OpType op_type(NodeId id) const {
        auto it = metadata_.find(id);
        PULSATRIX_ASSERT(it != metadata_.end());
        return it->second.op_type;
    }
 
    [[nodiscard]] std::optional<std::string> label(NodeId id) const {
        auto it = metadata_.find(id);
        PULSATRIX_ASSERT(it != metadata_.end());
        return it->second.label;
    }
 
private:
    std::unordered_map<NodeId, Tensor> activations_;
    std::vector<NodeId> node_ids_;
    std::unordered_map<NodeId, NodeMetadata> metadata_;
};
 
}  // namespace pulsatrix