| 55 | } |
| 56 | |
| 57 | EncodedSubgraph OpDef::make_backward_graph( |
| 58 | const OpDef& def, const SmallVector<LogicalTensorDesc>& inputs, |
| 59 | const SmallVector<bool>& input_requires_grad, |
| 60 | const SmallVector<bool>& output_has_grad) { |
| 61 | using BackwardGraphCache = |
| 62 | OpMethResultCache<EncodedSubgraph, SmallVector<bool>, SmallVector<bool>>; |
| 63 | thread_local auto& cache = *ResourceManager::create_local<BackwardGraphCache>(); |
| 64 | BackwardGraphCache::key_t cache_key{ |
| 65 | const_cast<OpDef&>(def).shared_from_this(), |
| 66 | inputs, |
| 67 | {input_requires_grad, output_has_grad}}; |
| 68 | auto iter = cache.find(cache_key); |
| 69 | if (iter == cache.end()) { |
| 70 | iter = cache.insert({cache_key, def.trait()->make_backward_graph( |
| 71 | def, inputs, input_requires_grad, |
| 72 | output_has_grad)}) |
| 73 | .first; |
| 74 | } |
| 75 | return iter->second; |
| 76 | } |
| 77 | |
| 78 | std::vector<std::pair<const char*, std::string>> OpDef::props(const OpDef& def) { |
| 79 | return def.trait()->props(def); |