parameters: (ys, xs, grads, cond_graph, body_graph, name, while_op,
                        maximum_iterations)
return_statement: return grad_func_graph, args
docstring: Builds and returns the gradient FuncGraph of `func_graph` and its args.
  
    The returned grad_func_graph must be called with the returned
    args + grad_func_graph.captures.
  
    Args:
      ys: A `Tensor` or list of tensors to be differentiated.
      xs: A `Tensor` or list of tensors to be used for differentiation.
      grads: The incoming grads for `ys`.
      cond_graph: FuncGraph for the forward cond function.
      body_graph: FuncGraph for the forward body function.
      name: Name of the returned gradient function.
      while_op: The forward While op.
      maximum_iterations: Tensor. The maximum number of iterations.
  
    Returns:
      2-tuple of (grad_func_graph, args).
docstring_summary: Builds and returns the gradient FuncGraph of `func_graph` and its args.
docstring_tokens: Builds and returns the gradient FuncGraph of func_graph and its args .
function: def _create_grad_func(ys, xs, grads, cond_graph, body_graph, name, while_op,
                        maximum_iterations):
    """Builds and returns the gradient FuncGraph of `func_graph` and its args.
  
    The returned grad_func_graph must be called with the returned
    args + grad_func_graph.captures.
  
    Args:
      ys: A `Tensor` or list of tensors to be differentiated.
      xs: A `Tensor` or list of tensors to be used for differentiation.
      grads: The incoming grads for `ys`.
      cond_graph: FuncGraph for the forward cond function.
      body_graph: FuncGraph for the forward body function.
      name: Name of the returned gradient function.
      while_op: The forward While op.
      maximum_iterations: Tensor. The maximum number of iterations.
  
    Returns:
      2-tuple of (grad_func_graph, args).
    """
    assert len(ys) == len(grads)
  
    total_iters = while_op.outputs[0]
    counter = constant_op.constant(
        0, dtype=total_iters.dtype, name="grad_counter")
  
    # Build frozen sets so that we do not have linear time lookups in
    # `_is_loop_invariant`. Note: `body_graph.inputs` and `body_graph.outputs`
    # may get updated during gradient computation because we add accumulators to
    # the forward op. However, those are not loop invariants so wouldn't affect
    # the output of `_is_loop_invariant`. Also we would never attempt to capture
    # those accumulators so `_is_loop_invariant` should never receive those new
    # tensors as args.
    body_graph_inputs = object_identity.ObjectIdentitySet(body_graph.inputs)
    body_graph_outputs = object_identity.ObjectIdentitySet(body_graph.outputs)
  
    args = [counter, maximum_iterations, total_iters] + list(grads)
    # Note: The returned function does not have `args` in the list of
    # `external_captures`.
    grad_func_graph = func_graph_module.func_graph_from_py_func(
        name,
        lambda *args: _grad_fn(ys, xs, args, body_graph),
        args, {},
        func_graph=_WhileBodyGradFuncGraph(name, cond_graph, body_graph,
                                           maximum_iterations, while_op,
                                           body_graph_inputs, body_graph_outputs))
  
    # Update the list of outputs with tensors corresponding to the captured
    # tensors. We capture 3 types of tensors when building the grad fn:
    # 1. Accumulators for forward graph intermediates which are not loop
    #    invariants. The outputs corresponding to these are populated in
    #    `popped_tensor_lists` by `_WhileBodyGradFuncGraph`.
    # 2. Resources, which are output as is.
    # 3. Forward graph loop invariants, which are output as is.
    for external_capture, internal_capture in grad_func_graph.captures:
      if ops.tensor_id(internal_capture) in grad_func_graph.popped_tensor_lists:
        new_output = grad_func_graph.popped_tensor_lists[ops.tensor_id(
            internal_capture)]
      elif (internal_capture.dtype == dtypes.resource or _is_loop_invariant(
          external_capture, body_graph_inputs, body_graph_outputs)):
        new_output = internal_capture
      else:
        raise ValueError("Tensor %s which captures %s is in list of "
                         "internal_captures but is not a resource, is not in "
                         "popped_tensor_lists and does not capture a loop "
                         "invariant." %
                         (str(internal_capture), str(external_capture)))
      grad_func_graph.outputs.append(new_output)
      grad_func_graph.structured_outputs.append(new_output)
  
    return grad_func_graph, args
function_tokens: def _create_grad_func ( ys , xs , grads , cond_graph , body_graph , name , while_op , maximum_iterations ) : assert len ( ys ) == len ( grads ) total_iters = while_op . outputs [ 0 ] counter = constant_op . constant ( 0 , dtype = total_iters . dtype , name = "grad_counter" ) # Build frozen sets so that we do not have linear time lookups in # `_is_loop_invariant`. Note: `body_graph.inputs` and `body_graph.outputs` # may get updated during gradient computation because we add accumulators to # the forward op. However, those are not loop invariants so wouldn't affect # the output of `_is_loop_invariant`. Also we would never attempt to capture # those accumulators so `_is_loop_invariant` should never receive those new # tensors as args. body_graph_inputs = object_identity . ObjectIdentitySet ( body_graph . inputs ) body_graph_outputs = object_identity . ObjectIdentitySet ( body_graph . outputs ) args = [ counter , maximum_iterations , total_iters ] + list ( grads ) # Note: The returned function does not have `args` in the list of # `external_captures`. grad_func_graph = func_graph_module . func_graph_from_py_func ( name , lambda * args : _grad_fn ( ys , xs , args , body_graph ) , args , { } , func_graph = _WhileBodyGradFuncGraph ( name , cond_graph , body_graph , maximum_iterations , while_op , body_graph_inputs , body_graph_outputs ) ) # Update the list of outputs with tensors corresponding to the captured # tensors. We capture 3 types of tensors when building the grad fn: # 1. Accumulators for forward graph intermediates which are not loop #    invariants. The outputs corresponding to these are populated in #    `popped_tensor_lists` by `_WhileBodyGradFuncGraph`. # 2. Resources, which are output as is. # 3. Forward graph loop invariants, which are output as is. for external_capture , internal_capture in grad_func_graph . captures : if ops . tensor_id ( internal_capture ) in grad_func_graph . popped_tensor_lists : new_output = grad_func_graph . popped_tensor_lists [ ops . tensor_id ( internal_capture ) ] elif ( internal_capture . dtype == dtypes . resource or _is_loop_invariant ( external_capture , body_graph_inputs , body_graph_outputs ) ) : new_output = internal_capture else : raise ValueError ( "Tensor %s which captures %s is in list of " "internal_captures but is not a resource, is not in " "popped_tensor_lists and does not capture a loop " "invariant." % ( str ( internal_capture ) , str ( external_capture ) ) ) grad_func_graph . outputs . append ( new_output ) grad_func_graph . structured_outputs . append ( new_output ) return grad_func_graph , args