| 435 | } |
| 436 | |
| 437 | Status SymbolicGradientBuilder::AddGradients() { |
| 438 | // Initialize backprops. |
| 439 | TF_RETURN_IF_ERROR(Initialize()); |
| 440 | |
| 441 | // Backward propagation. |
| 442 | std::vector<Output> dy; |
| 443 | while (!ready_.empty()) { |
| 444 | // n has collected all gradients. |
| 445 | Node* n = ready_.front(); |
| 446 | ready_.pop_front(); |
| 447 | |
| 448 | // dy[i] is the sum of i-th output's backpropped gradients. |
| 449 | const int num_y = n->num_outputs(); |
| 450 | dy.clear(); |
| 451 | dy.resize(num_y, {nullptr, 0}); |
| 452 | std::vector<int> no_grad_dy_indices; |
| 453 | for (int i = 0; i < num_y; ++i) { |
| 454 | TF_RETURN_IF_ERROR(SumGradients({n, i}, &dy[i])); |
| 455 | if (dy[i] == NoGradient()) { |
| 456 | no_grad_dy_indices.push_back(i); |
| 457 | } |
| 458 | auto iter = input_nodes_.find({n, i}); |
| 459 | if (iter != input_nodes_.end()) { |
| 460 | // Return gradients for Output in 'grad_outputs_'. |
| 461 | (*grad_outputs_)[iter->second] = dy[i]; |
| 462 | } |
| 463 | } |
| 464 | |
| 465 | // Stop backprop if none of the inputs to `n` are in `backprops_'. |
| 466 | bool stop_node = true; |
| 467 | for (const Edge* e : n->in_edges()) { |
| 468 | if (e->IsControlEdge()) continue; |
| 469 | if (backprops_.find({e->src(), e->src_output()}) != backprops_.end()) { |
| 470 | stop_node = false; |
| 471 | break; |
| 472 | } |
| 473 | } |
| 474 | |
| 475 | if (stop_node) { |
| 476 | continue; |
| 477 | } |
| 478 | |
| 479 | // Special case: if we find an exit node, process the associated while loop. |
| 480 | // Note that ProcessWhileLoop() calls BackpropAlongEdge() if necessary |
| 481 | // (which updates ready_), and we skip all the regular processing below |
| 482 | // after calling it. |
| 483 | if (n->IsExit()) { |
| 484 | DCHECK_EQ(dy.size(), 1); |
| 485 | TF_RETURN_IF_ERROR(ProcessWhileLoop(n, dy[0])); |
| 486 | continue; |
| 487 | } |
| 488 | // All loop-specific control flow ops should have been handled above |
| 489 | DCHECK(!n->IsEnter() && !n->IsNextIteration()) << n->DebugString(); |
| 490 | |
| 491 | const size_t num_no_grad = no_grad_dy_indices.size(); |
| 492 | if (IsPrimitiveOpWithNoGrad(n->type_string()) || num_no_grad == num_y) { |
| 493 | // No grad defined for this op, or all outputs returned 'NoGradient': |
| 494 | // Backprop 'NoGradient' along the in edges. |
no test coverage detected