Creates an If op from the specified predicate, branch functions and inputs. Note that this modifies true_graph and false_graph to make the inputs match, and to output all intermediates values so they're available for the gradient computation. true_graph and false_graph need not have the sa
(pred,
true_graph,
false_graph,
true_inputs,
false_inputs,
building_gradient,
name=None)
| 184 | |
| 185 | |
| 186 | def _build_cond(pred, |
| 187 | true_graph, |
| 188 | false_graph, |
| 189 | true_inputs, |
| 190 | false_inputs, |
| 191 | building_gradient, |
| 192 | name=None): |
| 193 | """Creates an If op from the specified predicate, branch functions and inputs. |
| 194 | |
| 195 | Note that this modifies true_graph and false_graph to make the inputs match, |
| 196 | and to output all intermediates values so they're available for the gradient |
| 197 | computation. |
| 198 | |
| 199 | true_graph and false_graph need not have the same input types, but they must |
| 200 | have the same outpute types. |
| 201 | |
| 202 | Args: |
| 203 | pred: boolean Tensor |
| 204 | true_graph: FuncGraph |
| 205 | false_graph: FuncGraph |
| 206 | true_inputs: a list of Tensors to be passed to true_graph as input. |
| 207 | false_inputs: a list of Tensors to be passed to false_graph as input. |
| 208 | building_gradient: Whether this is a gradient If op. |
| 209 | name: the name for the If op. |
| 210 | |
| 211 | Returns: |
| 212 | A list of Tensors which are the outputs of the If op. Does not include added |
| 213 | intermediate outputs. |
| 214 | """ |
| 215 | _make_indexed_slices_indices_types_match(_COND, [true_graph, false_graph]) |
| 216 | _check_same_outputs(_COND, [true_graph, false_graph]) |
| 217 | |
| 218 | # Add inputs to true_graph and false_graph to make them match. Note that |
| 219 | # this modifies true_graph and false_graph. |
| 220 | cond_inputs = _make_inputs_match([true_graph, false_graph], |
| 221 | [true_inputs, false_inputs]) |
| 222 | # Save the original number of outputs to return to the caller. |
| 223 | num_cond_outputs = len(true_graph.outputs) |
| 224 | # We do not output intermediates of the gradient If op since this is just |
| 225 | # for backwards compatibility with existing code. |
| 226 | if not building_gradient and util.output_all_intermediates(): |
| 227 | # Add all intermediate tensors as function outputs so they're available for |
| 228 | # the gradient computation. Since the outputs of the two functions must |
| 229 | # match, we wrap all the intermediates in optionals. Each intermediate |
| 230 | # output will have a value iff its corresponding branch is taken. |
| 231 | |
| 232 | true_intermediates = _get_intermediates(true_graph) |
| 233 | false_intermediates = _get_intermediates(false_graph) |
| 234 | |
| 235 | # Wrap intermediates in optionals. |
| 236 | wrapped_true_intermediates = _wrap_intermediates(true_graph, |
| 237 | true_intermediates) |
| 238 | wrapped_false_intermediates = _wrap_intermediates(false_graph, |
| 239 | false_intermediates) |
| 240 | |
| 241 | # Make outputs match by adding none optionals. |
| 242 | extra_true_outputs, extra_false_outputs = _make_intermediates_match( # pylint: disable=unbalanced-tuple-unpacking |
| 243 | [true_graph, false_graph], |
no test coverage detected