| 611 | } |
| 612 | |
| 613 | void GenEagerPythonOp::AddEagerFunctionTeardown( |
| 614 | const string& indentation, const std::vector<string>& output_sizes, |
| 615 | bool execute_record_gradient) { |
| 616 | if (num_outs_ > 0) { |
| 617 | if (execute_record_gradient) { |
| 618 | strings::StrAppend(&result_, indentation, "_execute.record_gradient(\n", |
| 619 | " \"", op_def_.name(), |
| 620 | "\", _inputs_flat, _attrs, _result, name)\n"); |
| 621 | } |
| 622 | if (num_outs_ == 1 && !output_sizes[0].empty()) { |
| 623 | // Single list result. |
| 624 | } else if (num_outs_ == 1) { |
| 625 | // Execute returns a single-element list which we need to destructure. |
| 626 | strings::StrAppend(&result_, indentation, "_result, = _result\n"); |
| 627 | } else { |
| 628 | // Have multiple outputs, so we will need to reformat the return |
| 629 | // value of execute() to be a list with one entry per op output |
| 630 | // (that entry will be a list of tensors if that output is of list |
| 631 | // type). |
| 632 | // For list outputs, convert the right subrange of _result into a list. |
| 633 | Unflatten(indentation, output_sizes, "_result", &result_); |
| 634 | // Convert to a named tuple. |
| 635 | strings::StrAppend(&result_, indentation, "_result = _", op_def_.name(), |
| 636 | "Output._make(_result)\n"); |
| 637 | } |
| 638 | } else { |
| 639 | strings::StrAppend(&result_, indentation, "_result = None\n"); |
| 640 | } |
| 641 | strings::StrAppend(&result_, indentation, "return _result\n\n"); |
| 642 | } |
| 643 | |
| 644 | bool GenEagerPythonOp::AddEagerFastPathAndGraphCode( |
| 645 | const string& parameters, const std::vector<string>& output_sizes, |