Helper class to manage calculating and caching control_outputs in graph.
| 66 | |
| 67 | |
| 68 | class _ControlOutputCache(object): |
| 69 | """Helper class to manage calculating and caching control_outputs in graph.""" |
| 70 | |
| 71 | def __init__(self): |
| 72 | self.cache = {} |
| 73 | |
| 74 | def calc_control_outputs(self, graph): |
| 75 | """Returns the map of control_outputs for a given graph. |
| 76 | |
| 77 | Args: |
| 78 | graph: The graph to parse. |
| 79 | |
| 80 | Returns: |
| 81 | A map of the control outputs. |
| 82 | """ |
| 83 | control_outputs = {} |
| 84 | for op in graph.get_operations(): |
| 85 | for control_input in op.control_inputs: |
| 86 | if control_input not in control_outputs: |
| 87 | control_outputs[control_input] = set() |
| 88 | control_outputs[control_input].add(op) |
| 89 | return control_outputs |
| 90 | |
| 91 | def get_control_outputs(self, op): |
| 92 | """Return the control outputs for a given op. |
| 93 | |
| 94 | Args: |
| 95 | op: The op to fetch control outputs for. |
| 96 | |
| 97 | Returns: |
| 98 | Iterable of control output ops. |
| 99 | """ |
| 100 | if op.graph not in self.cache: |
| 101 | control_outputs = self.calc_control_outputs(op.graph) |
| 102 | self.cache[op.graph] = control_outputs |
| 103 | else: |
| 104 | control_outputs = self.cache[op.graph] |
| 105 | return control_outputs.get(op, []) |
| 106 | |
| 107 | |
| 108 | def _subscribe_new(tensor, side_effects, control_cache): |