find_tape_or_recomp_stmts.h

Go to the documentation of this file.

#ifndef FREE_TENSOR_FIND_TAPE_OR_RECOMP_STMTS_H
#define FREE_TENSOR_FIND_TAPE_OR_RECOMP_STMTS_H
 
#include <unordered_set>
 
#include <autograd/derivative.h>
 
namespace freetensor {
 
std::pair<std::unordered_map<ID, std::unordered_set<ID>>,
          std::unordered_map<ID, std::unordered_set<ID>>>
findTapeOrRecompStmts(
    const Stmt &op, const std::unordered_set<ID> &defsToTape,
    const std::unordered_set<ID> defsNeedGrad,
    std::unordered_map<StmtOrExprID, Derivative::LazyFullDerivative>
        &derivatives);
 
} // namespace freetensor
 
#endif // FREE_TENSOR_FIND_TAPE_OR_RECOMP_STMTS_H