A `ControlFlowContext` for nodes inside an XLA computation cluster. THIS IS ONLY FOR TENSORFLOW INTERNAL IMPLEMENTATION, DO NO USE DIRECTLY. The primary role of `XLACompileContext` is to mark operators inside a xla.compile() computation with attribute "_xla_compile_id=XYZ", where XYZ is a
| 111 | |
| 112 | |
| 113 | class XLACompileContext(control_flow_ops.XLAControlFlowContext): |
| 114 | """A `ControlFlowContext` for nodes inside an XLA computation cluster. |
| 115 | |
| 116 | THIS IS ONLY FOR TENSORFLOW INTERNAL IMPLEMENTATION, DO NO USE DIRECTLY. |
| 117 | |
| 118 | The primary role of `XLACompileContext` is to mark operators inside a |
| 119 | xla.compile() computation with attribute "_xla_compile_id=XYZ", where XYZ is |
| 120 | a unique name. |
| 121 | |
| 122 | `ControlFlowContext` is used to perform the annotation since it integrates |
| 123 | with Tensorflow constructs like ResourceVariables. For example, if a |
| 124 | `ResourceVariable` is constructed inside a xla.compile() block, the |
| 125 | `ResourceVariable` implementation can use |
| 126 | `with ops.control_dependencies(None)` to build the variable's definition |
| 127 | outside the compiled computation. |
| 128 | """ |
| 129 | |
| 130 | def __init__(self, name, pivot): |
| 131 | """Builds a new XLACompileContext. |
| 132 | |
| 133 | Args: |
| 134 | name: a unique name for the context, used to populate the |
| 135 | `_xla_compile_id` attribute. |
| 136 | pivot: a pivot node. Nodes in the XLACompileContext that do not have any |
| 137 | inputs will have a control dependency on the pivot node. This ensures |
| 138 | that nodes are correctly included in any enclosing control flow |
| 139 | contexts. |
| 140 | """ |
| 141 | super(XLACompileContext, self).__init__() |
| 142 | self._name = name |
| 143 | self._name_as_bytes = compat.as_bytes(name) |
| 144 | self._unsupported_ops = [] |
| 145 | self._pivot = pivot |
| 146 | |
| 147 | def report_unsupported_operations(self): |
| 148 | if self._unsupported_ops: |
| 149 | op_str = '\n'.join([ |
| 150 | ' %s (%s)' % (op.type, op.name) |
| 151 | for op in self._unsupported_ops[:_MAX_WARNING_LINES] |
| 152 | ]) |
| 153 | logging.warning('%d unsupported operations found: \n%s', |
| 154 | len(self._unsupported_ops), op_str) |
| 155 | if len(self._unsupported_ops) > _MAX_WARNING_LINES: |
| 156 | logging.warning('... and %d more', |
| 157 | len(self._unsupported_ops) - _MAX_WARNING_LINES) |
| 158 | |
| 159 | def _RemoveExternalControlEdges(self, op): |
| 160 | """Remove any external control dependency on this op.""" |
| 161 | internal_control_inputs = [] |
| 162 | external_control_inputs = [] |
| 163 | for x in op.control_inputs: |
| 164 | # pylint: disable=protected-access |
| 165 | is_internal_op = False |
| 166 | ctxt = x._get_control_flow_context() |
| 167 | while ctxt is not None: |
| 168 | if ctxt == self: |
| 169 | is_internal_op = True |
| 170 | break |