Modifies each branch_graph's outputs to have the same output signature. Currently the only transformation implemented is turning a Tensor into an equivalent IndexedSlices if the other branch returns an IndexedSlices. Updates branch_graph.{outputs,structured_outputs} for each branch_graph in
(op_type, branch_graphs)
| 551 | |
| 552 | |
| 553 | def _make_output_composite_tensors_match(op_type, branch_graphs): |
| 554 | """Modifies each branch_graph's outputs to have the same output signature. |
| 555 | |
| 556 | Currently the only transformation implemented is turning a Tensor into an |
| 557 | equivalent IndexedSlices if the other branch returns an IndexedSlices. |
| 558 | Updates branch_graph.{outputs,structured_outputs} for each branch_graph in |
| 559 | branch_graphs. |
| 560 | |
| 561 | Args: |
| 562 | op_type: _COND or _CASE |
| 563 | branch_graphs: `list` of `FuncGraph` |
| 564 | |
| 565 | Raises: |
| 566 | TypeError: if a set of outputs cannot be rewritten. |
| 567 | """ |
| 568 | # Note: since this is only used for gradient graphs, we do not expect the |
| 569 | # outputs to be structured (e.g. nested lists), and thus do not need to use |
| 570 | # nest.flatten, etc. |
| 571 | assert branch_graphs |
| 572 | branch_outputs = [g.structured_outputs for g in branch_graphs] |
| 573 | outputs_per_branch = list(len(outs) for outs in branch_outputs) |
| 574 | assert len(set(outputs_per_branch)) == 1, outputs_per_branch |
| 575 | |
| 576 | for output_idx, branch_outs in enumerate(zip(*branch_outputs)): |
| 577 | if len(set(type(out) for out in branch_outs)) == 1: |
| 578 | continue |
| 579 | if not any(isinstance(out, ops.IndexedSlices) for out in branch_outs): |
| 580 | continue |
| 581 | for branch_idx, branch_out in enumerate(branch_outs): |
| 582 | if isinstance(branch_out, ops.IndexedSlices): |
| 583 | continue |
| 584 | elif isinstance(branch_out, ops.Tensor): |
| 585 | with branch_graphs[branch_idx].as_default(): |
| 586 | branch_outputs[branch_idx][output_idx] = math_ops._as_indexed_slices( |
| 587 | branch_out) |
| 588 | else: |
| 589 | raise TypeError( |
| 590 | "Cannot reconcile {op_name} {output_idx}-th outputs:\n" |
| 591 | " outputs from all branches: {outputs}".format( |
| 592 | op_name="tf.cond" if op_type == _COND else "tf.switch_case", |
| 593 | output_idx=output_idx, |
| 594 | outputs=branch_outs)) |
| 595 | |
| 596 | for branch_graph, branch_outs in zip(branch_graphs, branch_outputs): |
| 597 | branch_graph.structured_outputs = branch_outs |
| 598 | branch_graph.outputs = func_graph_module.flatten(branch_outs) |
| 599 | |
| 600 | |
| 601 | def _make_indexed_slices_indices_types_match(op_type, branch_graphs): |