FuncGraph for the gradient function of the branch of an If op. Handles wrapping and unwrapping intermediate values that are captured by the gradient computation in optionals. Attributes: op_needs_rewrite: True if any intermediates were captured, meaning the forward If op needs to b
| 769 | |
| 770 | |
| 771 | class _CondGradFuncGraph(util.CondBranchFuncGraph): |
| 772 | """FuncGraph for the gradient function of the branch of an If op. |
| 773 | |
| 774 | Handles wrapping and unwrapping intermediate values that are captured by the |
| 775 | gradient computation in optionals. |
| 776 | |
| 777 | Attributes: |
| 778 | op_needs_rewrite: True if any intermediates were captured, meaning the |
| 779 | forward If op needs to be written to output the wrapped intermediates. |
| 780 | """ |
| 781 | |
| 782 | def __init__(self, name, forward_graph): |
| 783 | super(_CondGradFuncGraph, self).__init__( |
| 784 | name, collections=ops.get_default_graph()._collections) # pylint: disable=protected-access |
| 785 | self.op_needs_rewrite = False |
| 786 | self._forward_graph = forward_graph |
| 787 | # Maps from forward intermediate tensor -> the unwrapped captured |
| 788 | # intermediate. |
| 789 | self._indirect_captures = {} |
| 790 | # Maps unwrapped intermediate -> optional-wrapped intermediate in the |
| 791 | # forward graph. |
| 792 | self._wrapped_intermediates = collections.OrderedDict() |
| 793 | # Raw intermediates captured from the forward graph. Populated iff we're in |
| 794 | # an XLA context. |
| 795 | self._xla_intermediates = [] |
| 796 | |
| 797 | @property |
| 798 | def wrapped_intermediates(self): |
| 799 | """The optional-wrapped intermediates captured from the forward graph.""" |
| 800 | return list(self._wrapped_intermediates.values()) |
| 801 | |
| 802 | @property |
| 803 | def xla_intermediates(self): |
| 804 | """Raw intermediates captured from the forward graph if XLA is enabled.""" |
| 805 | return self._xla_intermediates |
| 806 | |
| 807 | def _capture_helper(self, tensor, name): |
| 808 | if (tensor.graph is not self._forward_graph or |
| 809 | any(tensor is t for t in self._forward_graph.inputs) or |
| 810 | any(tensor is t for t in self._forward_graph.outputs)): |
| 811 | return super(_CondGradFuncGraph, self)._capture_helper(tensor, name) |
| 812 | |
| 813 | if control_flow_util.GraphOrParentsInXlaContext(ops.get_default_graph()): |
| 814 | # XLA does not yet support optionals, so capture intermediates directly. |
| 815 | # TODO(skyewm,jpienaar): can XLA support optionals? |
| 816 | if all(tensor is not capture for capture in self.external_captures): |
| 817 | self.xla_intermediates.append(tensor) |
| 818 | self.op_needs_rewrite = True |
| 819 | return super(_CondGradFuncGraph, self)._capture_helper(tensor, name) |
| 820 | |
| 821 | tensor_id = ops.tensor_id(tensor) |
| 822 | captured_tensor = self._indirect_captures.get(tensor_id) |
| 823 | if captured_tensor is not None: |
| 824 | return captured_tensor |
| 825 | |
| 826 | # 'tensor' is an uncaptured intermediate in the forward graph. |
| 827 | # If it is not a resource, we wrap it in an optional in the forward graph |
| 828 | # and capture the optional normally. We then unwrap the captured optional |