(self, tensor, name)
| 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 |
| 829 | # value in the gradient graph to get the raw intermediate value. |
| 830 | # If it is a resource, we trace the resource upto the input in the forward |
| 831 | # graph and capture that. |
| 832 | |
| 833 | if tensor.dtype == dtypes.resource: |
| 834 | # Index of the forward graph input corresponding to the resource tensor. |
| 835 | index = util.resource_input_index( |
| 836 | tensor.name, [t.name for t in self._forward_graph.inputs], |
| 837 | {op.name: op.node_def for op in self._forward_graph.get_operations()}, |
| 838 | self._forward_graph._functions) |
| 839 | # This gets mapped to the corresponding If op input in |
| 840 | # `_resolve_grad_inputs`. |
| 841 | captured_tensor = super(_CondGradFuncGraph, self)._capture_helper( |
| 842 | self._forward_graph.inputs[index], name) |
| 843 | else: |
| 844 | if tensor_id not in self._wrapped_intermediates: |
| 845 | # If the gradient has already been computed for this If op, 'tensor' may |
| 846 | # already be wrapped. |
| 847 | for consumer in tensor.consumers(): |
| 848 | if (consumer.type == "OptionalFromValue" and |
| 849 | any(consumer.outputs[0] is output |
| 850 | for output in self._forward_graph.outputs)): |
| 851 | optional = consumer.outputs[0] |
| 852 | break |
| 853 | else: |
| 854 | # 'tensor' hasn't been wrapped, do it now. |
| 855 | with self._forward_graph.as_default(): |
| 856 | optional = gen_dataset_ops.optional_from_value([tensor]) |
| 857 | self.op_needs_rewrite = True |
| 858 | self._wrapped_intermediates[tensor_id] = optional |
| 859 | |
| 860 | optional = self._wrapped_intermediates[tensor_id] |
| 861 | captured_optional = super(_CondGradFuncGraph, |
| 862 | self)._capture_helper(optional, name) |
| 863 | captured_tensor = gen_dataset_ops.optional_get_value( |
| 864 | captured_optional, [tensor.dtype], [tensor.shape])[0] |
nothing calls this directly
no test coverage detected