| 99 | } |
| 100 | |
| 101 | static Node* AddSymGrad(Graph* g, Node* n, gtl::ArraySlice<NodeOut> grads) { |
| 102 | const int num_x = n->num_inputs(); |
| 103 | const int num_y = n->num_outputs(); |
| 104 | CHECK_EQ(num_y, grads.size()); |
| 105 | |
| 106 | NodeDef ndef; |
| 107 | ndef.set_name(g->NewName(kNodeLabel)); |
| 108 | ndef.set_op(kGradientOp); |
| 109 | |
| 110 | // The gradient node should have num_x + num_y inputs. |
| 111 | std::vector<NodeOut> n_inputs(num_x); |
| 112 | for (const Edge* e : n->in_edges()) { |
| 113 | if (e->IsControlEdge()) continue; |
| 114 | n_inputs[e->dst_input()] = {e->src(), e->src_output()}; |
| 115 | } |
| 116 | DataTypeVector in_types; |
| 117 | for (const NodeOut& nout : n_inputs) { |
| 118 | ndef.add_input(nout.name()); |
| 119 | in_types.push_back(nout.dtype()); |
| 120 | } |
| 121 | for (const NodeOut& nout : grads) { |
| 122 | ndef.add_input(nout.name()); |
| 123 | in_types.push_back(nout.dtype()); |
| 124 | } |
| 125 | CHECK_EQ(ndef.input_size(), num_x + num_y); |
| 126 | |
| 127 | AddNodeAttr("Tin", in_types, &ndef); |
| 128 | |
| 129 | // The gradient node's outputs have the same types as the node 'n's |
| 130 | // inputs, except for resources. |
| 131 | DataTypeVector out_types = n->input_types(); |
| 132 | for (int i = 0; i < out_types.size(); ++i) { |
| 133 | if (out_types[i] == DT_RESOURCE) { |
| 134 | // TODO(apassos): figure out how to get the right dtype |
| 135 | out_types[i] = DT_FLOAT; |
| 136 | } |
| 137 | } |
| 138 | AddNodeAttr("Tout", out_types, &ndef); |
| 139 | NameAttrList func; |
| 140 | func.set_name(n->type_string()); |
| 141 | for (const auto& attr : n->attrs()) { |
| 142 | (*func.mutable_attr())[attr.first] = attr.second; |
| 143 | } |
| 144 | AddNodeAttr("f", func, &ndef); |
| 145 | Status s; |
| 146 | Node* ret = g->AddNode(ndef, &s); |
| 147 | TF_CHECK_OK(s); |
| 148 | return ret; |
| 149 | } |
| 150 | |
| 151 | class SymbolicGradientBuilder { |
| 152 | public: |
no test coverage detected