A `ControlFlowContext` for nodes inside a TPU inference computation. The primary role of `TPUReplicateContext` is to sanity check operators inside a tpu.rewrite_for_inference() computation.
| 1544 | |
| 1545 | |
| 1546 | class _TPUInferenceContext(control_flow_ops.XLAControlFlowContext): |
| 1547 | """A `ControlFlowContext` for nodes inside a TPU inference computation. |
| 1548 | |
| 1549 | The primary role of `TPUReplicateContext` is to sanity check operators inside |
| 1550 | a tpu.rewrite_for_inference() computation. |
| 1551 | """ |
| 1552 | |
| 1553 | def __init__(self, name): |
| 1554 | super(_TPUInferenceContext, self).__init__() |
| 1555 | self._name = name |
| 1556 | |
| 1557 | def AddOp(self, op): |
| 1558 | self._AddOpInternal(op) |
| 1559 | |
| 1560 | def _AddOpInternal(self, op): |
| 1561 | # pylint: disable=protected-access |
| 1562 | if op.type in _BLACKLISTED_INFERENCE_OPS: |
| 1563 | raise NotImplementedError( |
| 1564 | "Operation of type %s (%s) is not supported on the TPU for inference." |
| 1565 | " Execution will fail if this op is used in the graph. Make sure your" |
| 1566 | " variables are using variable_scope." % (op.type, op.name)) |
| 1567 | if self._outer_context: |
| 1568 | self._outer_context.AddInnerOp(op) |
| 1569 | |
| 1570 | def AddValue(self, val): |
| 1571 | result = val |
| 1572 | if self._outer_context: |
| 1573 | result = self._outer_context.AddValue(val) |
| 1574 | return result |
| 1575 | |
| 1576 | def AddInnerOp(self, op): |
| 1577 | self._AddOpInternal(op) |
| 1578 | |
| 1579 | @property |
| 1580 | def grad_state(self): |
| 1581 | return None |
| 1582 | |
| 1583 | |
| 1584 | def validate_inference_rewrite_for_variables(graph): |
no outgoing calls
no test coverage detected