autograd.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <functional>
#include <unordered_map>
#include <utility>
#include <vector>
 
#include "pulsatrix/computation_graph.hpp"
#include "pulsatrix/node.hpp"
#include "pulsatrix/tensor.hpp"
 
namespace pulsatrix {
 
class Autograd {
public:
    using BackwardFn = std::function<Tensor(const Tensor& grad_output)>;
 
    void register_backward(NodeId id, BackwardFn fn);
 
    void backward(const ComputationGraph& graph, NodeId root, const Tensor& grad_output);
 
    void backward(const ComputationGraph& graph, std::vector<std::pair<NodeId, Tensor>> seeds);
 
    [[nodiscard]] const Tensor& gradient(NodeId id) const;
 
    [[nodiscard]] bool has_gradient(NodeId id) const { return gradients_.find(id) != gradients_.end(); }
 
private:
    std::unordered_map<NodeId, BackwardFn> backward_fns_;
    std::unordered_map<NodeId, Tensor> gradients_;
};
 
}  // namespace pulsatrix