| 475 | // are needed, are copied and returned in BackpropInitialState. |
| 476 | template <typename BackwardFunction, typename TapeTensor> |
| 477 | BackpropInitialState<BackwardFunction, TapeTensor> PrepareBackprop( |
| 478 | gtl::ArraySlice<int64> target, const TensorTape& tensor_tape, |
| 479 | OpTape<BackwardFunction, TapeTensor>* op_tape, |
| 480 | const std::unordered_set<int64>& sources_set, bool persistent_tape) { |
| 481 | std::vector<int64> tensor_stack; |
| 482 | tensor_stack.reserve(target.size()); |
| 483 | for (auto t : target) { |
| 484 | tensor_stack.push_back(t); |
| 485 | } |
| 486 | BackpropInitialState<BackwardFunction, TapeTensor> result; |
| 487 | while (!tensor_stack.empty()) { |
| 488 | int64 tensor_id = tensor_stack.back(); |
| 489 | tensor_stack.pop_back(); |
| 490 | auto op_id_it = tensor_tape.find(tensor_id); |
| 491 | if (op_id_it == tensor_tape.end()) { |
| 492 | continue; |
| 493 | } |
| 494 | int64 op_id = op_id_it->second; |
| 495 | auto op_it = op_tape->find(op_id); |
| 496 | auto result_op_it = result.op_tape.find(op_id); |
| 497 | if (op_id == -1 || op_it == op_tape->end() || |
| 498 | result_op_it != result.op_tape.end()) { |
| 499 | continue; |
| 500 | } |
| 501 | CHECK(result.op_tape.emplace(op_id, op_it->second).second); |
| 502 | for (auto it : op_it->second.input_tensor_id) { |
| 503 | auto count_it = result.tensor_usage_counts.find(it); |
| 504 | if (count_it != result.tensor_usage_counts.end()) { |
| 505 | count_it->second++; |
| 506 | } else { |
| 507 | result.tensor_usage_counts[it] = 1; |
| 508 | if (tensor_tape.find(it) != tensor_tape.end()) { |
| 509 | tensor_stack.push_back(it); |
| 510 | } |
| 511 | } |
| 512 | } |
| 513 | if (!persistent_tape) { |
| 514 | op_tape->erase(op_it); |
| 515 | } |
| 516 | } |
| 517 | for (auto& pair : result.tensor_usage_counts) { |
| 518 | auto it = tensor_tape.find(pair.first); |
| 519 | if (it != tensor_tape.end() && it->second != -1) { |
| 520 | result.op_missing_tensor[it->second] += 1; |
| 521 | } |
| 522 | } |
| 523 | if (!persistent_tape) { |
| 524 | // Call destructors for all unneeded gradient functions and |
| 525 | // clear the op_tape. We can clear the tape because ownership of |
| 526 | // backward functions that will be used for gradient computation |
| 527 | // has been transferred to `result`. |
| 528 | for (const auto& op_pair : *op_tape) { |
| 529 | op_pair.second.backward_function_deleter( |
| 530 | op_pair.second.backward_function); |
| 531 | } |
| 532 | op_tape->clear(); |
| 533 | } |
| 534 | return result; |