Execute computation under `_TPUInferenceContext`.
(*args, **kwargs)
| 1649 | getter(name, *args, **kwargs), name=name + "/GuaranteeConst") |
| 1650 | |
| 1651 | def wrapped_computation(*args, **kwargs): |
| 1652 | """Execute computation under `_TPUInferenceContext`.""" |
| 1653 | context = _TPUInferenceContext( |
| 1654 | name=ops.get_default_graph().unique_name("rewrite_for_inference")) |
| 1655 | try: |
| 1656 | context.Enter() |
| 1657 | |
| 1658 | vscope = variable_scope.get_variable_scope() |
| 1659 | prev_custom_getter = vscope.custom_getter |
| 1660 | prev_caching_device = vscope.caching_device |
| 1661 | vscope.set_custom_getter(guarantee_const_getter) |
| 1662 | vscope.set_caching_device(lambda op: op.device) |
| 1663 | |
| 1664 | result = computation(*args, **kwargs) |
| 1665 | |
| 1666 | vscope.set_custom_getter(prev_custom_getter) |
| 1667 | vscope.set_caching_device(prev_caching_device) |
| 1668 | finally: |
| 1669 | context.Exit() |
| 1670 | return result |
| 1671 | |
| 1672 | # pylint: disable=undefined-variable |
| 1673 | return rewrite( |
nothing calls this directly
no test coverage detected