Sets the flag to enable lowering on `op` if necessary. Lowering allows cond_v2 and while_v2 to avoid some of the limitations of Functions, allowing users to specify devices & colocation inside of cond_v2 and while_v2 input functions, and enabling non-strict evaluation & partial pruning. Thi
(op)
| 81 | |
| 82 | |
| 83 | def maybe_set_lowering_attr(op): |
| 84 | """Sets the flag to enable lowering on `op` if necessary. |
| 85 | |
| 86 | Lowering allows cond_v2 and while_v2 to avoid some of the limitations of |
| 87 | Functions, allowing users to specify devices & colocation inside of cond_v2 |
| 88 | and while_v2 input functions, and enabling non-strict evaluation & partial |
| 89 | pruning. This brings v2 control flow closer to feature parity with v1 control |
| 90 | flow. |
| 91 | |
| 92 | However, we do not lower in the following cases: |
| 93 | - When the `If` or `While` ops are in the XLA context. Because it is easier |
| 94 | for XLA to apply its own optimizations when dealing with un-lowered |
| 95 | control flow operators than with low-level control flow primitives. |
| 96 | - When the eager execution context specifies the executor of functions to |
| 97 | be the single threaded executor (see context.function_executor_type()). |
| 98 | Because the single threaded executor does not support v1 control flow ops. |
| 99 | |
| 100 | Args: |
| 101 | op: An `If` or `While` Operation. |
| 102 | """ |
| 103 | if (not control_flow_util.GraphOrParentsInXlaContext(op.graph) and |
| 104 | context.context().function_call_options.executor_type != |
| 105 | "SINGLE_THREADED_EXECUTOR"): |
| 106 | # pylint: disable=protected-access |
| 107 | op._set_attr("_lower_using_switch_merge", attr_value_pb2.AttrValue(b=True)) |
| 108 | # pylint: enable=protected-access |
| 109 | |
| 110 | |
| 111 | def maybe_propagate_compile_time_consts_in_xla(op): |