message: Minor refactor of _WhileBodyGradFuncGraph._capture_helper.
  It was getting kind of long. This also makes minor stylistic changes to the moved code.
diff: @@ -724,31 +724,9 @@ class _WhileBodyGradFuncGraph(util.WhileBodyFuncGraph):
  if captured_tensor is not None:
  return captured_tensor
  + # Resource tensors are not accumulated and handled specially.
  if tensor.dtype == dtypes.resource:
  - # Resource-type tensors are not accumulated.
  - # If a resource tensor exists in the loop body it must either be a loop
  - # input or an output of a nested While op inside the loop body which
  - # had captured the external resource.
  - if tensor in self._forward_graph.inputs:
  - index = self._forward_graph.inputs.index(tensor)
  - elif tensor.op.type == "While":
  - # Captured resources occur at the same index in the lists of inputs and
  - # outputs of a while op. So we lookup the input of `tensor.op` at the
  - # same index as the index of `tensor` in the `tensor.op.outputs`.
  - index = self._forward_graph.inputs.index(
  - tensor.op.inputs[tensor.value_index])
  - else:
  - raise ValueError(
  - "Taking gradient of a while loop which creates"
  - " a resource in its body is not supported: %s" % str(tensor))
  - # This must be a loop invariant.
  - assert self._forward_graph.inputs[index] == self._forward_graph.outputs[
  - index], "Resource tensors must be loop invariants %s." % str(
  - self._forward_graph._while.inputs[index])
  - tensor_in_outer_graph = self._forward_graph._while.inputs[index]
  - self._indirect_captures[tensor] = self.capture(
  - tensor_in_outer_graph, whitelisted=True)
  - return self._indirect_captures[tensor]
  + return self._resource_capture_helper(tensor)
  # Create or find an existing accumulator output for `tensor` in the forward
  # graph, and fetch from this accumulator in the gradient graph to get the
  @@ -789,6 +767,41 @@ class _WhileBodyGradFuncGraph(util.WhileBodyFuncGraph):
  self.popped_tensor_lists[captured_accumulator] = new_tensor_list
  return captured_tensor
  + def _resource_capture_helper(self, tensor):
  + """Returns the captured resource tensor.
  +
  + Resource-type tensors are not accumulated. If a resource tensor exists in
  + the loop body it must either be a loop input or an output of a nested While
  + op inside the loop body which had captured the external resource.
  +
  + Args:
  + tensor: the external resource Tensor to be captured.
  +
  + Returns:
  + Tensor in this graph.
  + """
  + assert tensor.dtype == dtypes.resource
  + if tensor in self._forward_graph.inputs:
  + index = self._forward_graph.inputs.index(tensor)
  + elif tensor.op.type == "While":
  + # Captured resources occur at the same index in the lists of inputs and
  + # outputs of a while op. So we lookup the input of `tensor.op` at the
  + # same index as the index of `tensor` in the `tensor.op.outputs`.
  + index = self._forward_graph.inputs.index(
  + tensor.op.inputs[tensor.value_index])
  + else:
  + raise ValueError(
  + "Taking gradient of a while loop which creates "
  + "a resource in its body is not supported: %s" % tensor)
  + # This must be a loop invariant.
  + assert self._forward_graph.inputs[index] == self._forward_graph.outputs[
  + index], ("Resource tensors must be loop invariants %s." %
  + self._forward_graph._while.inputs[index])
  + tensor_in_outer_graph = self._forward_graph._while.inputs[index]
  + self._indirect_captures[tensor] = self.capture(
  + tensor_in_outer_graph, whitelisted=True)
  + return self._indirect_captures[tensor]
  +
  def _check_shapes_compat(output_tensors, shape_invariants, input_tensors):
  for (t, shape, input_t) in zip(output_tensors, shape_invariants,
  