gradtable.h

Go to the documentation of this file.

#pragma once

#include <map>
#include <string>
#include <cstdint>
#include "compute/operation.h"
#include "compute/variable.h"

namespace magmadnn {
namespace op {

template <typename T>
class GradTable {
public:
    GradTable();

    unsigned int get_size();

    Tensor<T>* get(Operation<T>* var);

    void set(Operation<T>* var, Tensor<T>* grad);

    void clear();

protected:
    std::map<Operation<T> *, Tensor<T>* > _table;   // the underlying table to store data
    typename std::map<Operation<T> *, Tensor<T> *>::iterator tmp_map_iterator;

};

}   // namespace op
}   // namespace magmadnn