| 1696 | |
| 1697 | |
| 1698 | void AggregateForSparseApply(Graph* dest, Node* n, |
| 1699 | map_node_list& duplicated_nodes, int grad_input, |
| 1700 | int indices_input) { |
| 1701 | const Edge* e = nullptr; |
| 1702 | TF_CHECK_OK(n->input_edge(grad_input, &e)); |
| 1703 | auto grad = e->src(); |
| 1704 | auto dup_grads = duplicated_nodes[grad]; |
| 1705 | |
| 1706 | const Edge* e2 = nullptr; |
| 1707 | TF_CHECK_OK(n->input_edge(indices_input, &e2)); |
| 1708 | auto indices = e2->src(); |
| 1709 | auto dup_indices = duplicated_nodes[indices]; |
| 1710 | |
| 1711 | Node* axis = scalar_const_node(dest, 0, n->name() + "/axis_0", |
| 1712 | n->assigned_device_name()); |
| 1713 | |
| 1714 | std::vector<NodeDefBuilder::NodeOut> src_list; |
| 1715 | src_list.emplace_back(grad->name(), e->src_output(), |
| 1716 | grad->output_type(0)); |
| 1717 | for (auto it : dup_grads) { |
| 1718 | src_list.emplace_back(it->name(), e->src_output(), |
| 1719 | it->output_type(0)); |
| 1720 | } |
| 1721 | NodeDef node_def; |
| 1722 | TF_CHECK_OK(NodeDefBuilder(n->name() + "/aggr_concat", "ConcatV2") |
| 1723 | .Device(n->assigned_device_name()) |
| 1724 | .Input(src_list) |
| 1725 | .Input(axis->name(), 0, axis->output_type(0)) |
| 1726 | .Attr("T", grad->output_type(0)) |
| 1727 | .Finalize(&node_def)); |
| 1728 | Status s; |
| 1729 | Node* concat_grad = dest->AddNode(node_def, &s); |
| 1730 | TF_CHECK_OK(s); |
| 1731 | concat_grad->set_assigned_device_name(n->assigned_device_name()); |
| 1732 | |
| 1733 | std::vector<NodeDefBuilder::NodeOut> src_list2; |
| 1734 | src_list2.emplace_back(indices->name(), e2->src_output(), |
| 1735 | indices->output_type(0)); |
| 1736 | for (auto it : dup_indices) { |
| 1737 | src_list2.emplace_back(it->name(), e2->src_output(), |
| 1738 | it->output_type(0)); |
| 1739 | } |
| 1740 | NodeDef node_def2; |
| 1741 | TF_CHECK_OK(NodeDefBuilder(n->name() + "/aggr_concat2", "ConcatV2") |
| 1742 | .Device(n->assigned_device_name()) |
| 1743 | .Input(src_list2) |
| 1744 | .Input(axis->name(), 0, axis->output_type(0)) |
| 1745 | .Attr("T", indices->output_type(0)) |
| 1746 | .Finalize(&node_def2)); |
| 1747 | Node* concat_indices = dest->AddNode(node_def2, &s); |
| 1748 | TF_CHECK_OK(s); |
| 1749 | concat_indices->set_assigned_device_name(n->assigned_device_name()); |
| 1750 | |
| 1751 | dest->RemoveEdge(e); |
| 1752 | dest->RemoveEdge(e2); |
| 1753 | dest->AddEdge(grad, 0, concat_grad, 0); |
| 1754 | for (size_t i = 0; i < dup_grads.size(); ++i) { |
| 1755 | dest->AddEdge(dup_grads[i], 0, concat_grad, i+1); |
no test coverage detected