| 549 | |
| 550 | template <typename Gradient, typename BackwardFunction, typename TapeTensor> |
| 551 | Status InitialGradients( |
| 552 | const VSpace<Gradient, BackwardFunction, TapeTensor>& vspace, |
| 553 | gtl::ArraySlice<int64> target_tensor_ids, |
| 554 | const std::unordered_map<int64, TapeTensor>& sources_that_are_targets, |
| 555 | gtl::ArraySlice<Gradient*> output_gradients, const TensorTape& tensor_tape, |
| 556 | const OpTape<BackwardFunction, TapeTensor>& op_tape, |
| 557 | std::unordered_map<int64, std::vector<Gradient*>>* result) { |
| 558 | for (int i = 0; i < target_tensor_ids.size(); ++i) { |
| 559 | const int64 id = target_tensor_ids[i]; |
| 560 | if (output_gradients.empty() || output_gradients[i] == nullptr) { |
| 561 | auto tensor_it = tensor_tape.find(id); |
| 562 | if (tensor_it != tensor_tape.end() && tensor_it->second != -1) { |
| 563 | auto op_it = op_tape.find(tensor_it->second); |
| 564 | if (op_it == op_tape.end()) { |
| 565 | return errors::Internal( |
| 566 | "Internal state of the gradient tape is invalid: " |
| 567 | "failed to find operation producing a tensor"); |
| 568 | } |
| 569 | bool found = false; |
| 570 | for (int j = 0; j < op_it->second.output_tensor_info.size(); ++j) { |
| 571 | if (op_it->second.output_tensor_info[j].GetID() == id) { |
| 572 | found = true; |
| 573 | (*result)[id].push_back( |
| 574 | vspace.Ones(op_it->second.output_tensor_info[j])); |
| 575 | break; |
| 576 | } |
| 577 | } |
| 578 | if (!found) { |
| 579 | return errors::Internal( |
| 580 | "Internal state of the gradient tape is invalid: " |
| 581 | "none of operations outputs match expected tensor"); |
| 582 | } |
| 583 | } else { |
| 584 | // This target tensor was not generated by any operation recorded on |
| 585 | // the tape, so no gradient needs to be computed from it unless this |
| 586 | // target is also a source. |
| 587 | auto source_tensor = sources_that_are_targets.find(id); |
| 588 | if (source_tensor != sources_that_are_targets.end()) { |
| 589 | (*result)[id].push_back(vspace.Ones(source_tensor->second)); |
| 590 | } |
| 591 | } |
| 592 | } else { |
| 593 | (*result)[id].push_back(output_gradients[i]); |
| 594 | } |
| 595 | } |
| 596 | return Status::OK(); |
| 597 | } |
| 598 | |
| 599 | // TODO(agarwal): use an automatic mechanism for handling None arguments to |
| 600 | // gradient functions. |