Creates an `Case` op from `branch_index`, branch graphs and inputs. Note that this modifies `branch_graphs` to make the inputs match, and to output all intermediates values so they're available for the gradient computation. `branch_graphs` need not have the same input types, but they must
(branch_index, branch_graphs, branch_inputs, name=None)
| 987 | |
| 988 | |
| 989 | def _build_case(branch_index, branch_graphs, branch_inputs, name=None): |
| 990 | """Creates an `Case` op from `branch_index`, branch graphs and inputs. |
| 991 | |
| 992 | Note that this modifies `branch_graphs` to make the inputs match, and to |
| 993 | output all intermediates values so they're available for the gradient |
| 994 | computation. |
| 995 | |
| 996 | `branch_graphs` need not have the same input types, but they must |
| 997 | have the same outpute types. |
| 998 | |
| 999 | Args: |
| 1000 | branch_index: integer Tensor |
| 1001 | branch_graphs: List of FuncGraph |
| 1002 | branch_inputs: List of lists of Tensors to be passed to corresponding |
| 1003 | branch_graph as input. |
| 1004 | name: the name for the Case op. |
| 1005 | |
| 1006 | Returns: |
| 1007 | A list of Tensors which are the outputs of the Case op. Does not include |
| 1008 | added intermediate outputs. |
| 1009 | """ |
| 1010 | _make_indexed_slices_indices_types_match(_CASE, branch_graphs) |
| 1011 | _check_same_outputs(_CASE, branch_graphs) |
| 1012 | |
| 1013 | # Add inputs to branch_graphs to make them match. Note that this modifies the |
| 1014 | # graphs in `branch_graphs`. |
| 1015 | case_inputs = _make_inputs_match(branch_graphs, branch_inputs) |
| 1016 | |
| 1017 | # Create the Case op. |
| 1018 | with ops.control_dependencies( |
| 1019 | sum((list(bg.control_captures) for bg in branch_graphs), [])): |
| 1020 | tensors = gen_functional_ops.case( |
| 1021 | branch_index, |
| 1022 | case_inputs, [t.dtype for t in branch_graphs[0].outputs], |
| 1023 | [util.create_new_tf_function(g) for g in branch_graphs], |
| 1024 | output_shapes=_get_output_shapes(*[g.outputs for g in branch_graphs]), |
| 1025 | name=name) |
| 1026 | |
| 1027 | # TODO(b/110167197): this requires Case to have at least 1 output |
| 1028 | case_op = tensors[0].op |
| 1029 | util.maybe_set_lowering_attr(case_op) |
| 1030 | util.maybe_propagate_compile_time_consts_in_xla(case_op) |
| 1031 | |
| 1032 | # Return identities for each output of the Case op, rather than the output of |
| 1033 | # the Case op directly. This makes pruning work if the output of switch_case() |
| 1034 | # is fetched: the lowering pass converts the Case outputs into IdentityN |
| 1035 | # outputs, which if fetched will cause all ops in the taken branch to be run |
| 1036 | # (since it takes all merge ops as input). After lowering, each output |
| 1037 | # identity op will end up with only the appropriate merge op as input. |
| 1038 | # TODO(b/79984175): this doesn't have to be a tuple once we covert to the |
| 1039 | # correct output structure |
| 1040 | tensors = [array_ops.identity(t) for t in tensors] |
| 1041 | |
| 1042 | # Prevent fetching since the variant outputs can't be fetched directly. |
| 1043 | case_op.graph.prevent_fetching(case_op) |
| 1044 | return func_graph_module.pack_sequence_as(branch_graphs[0].structured_outputs, |
| 1045 | tensors) |
no test coverage detected