parameters: (op, *grads)
return_statement: return _get_structured_grad_output(outputs, grads, body_grad_graph)
docstring: The gradient of a While op produced by while_loop.
docstring_summary: The gradient of a While op produced by while_loop.
docstring_tokens: The gradient of a While op produced by while_loop .
function: def _WhileGrad(op, *grads):  # pylint: disable=invalid-name
    """The gradient of a While op produced by while_loop."""
    # Note that op is not always the same as while_op because the gradient tape,
    # for eager mode compatibility, forgets information about the proper op. Since
    # the loop cannot run in eager mode, however, we can safely introspect into
    # the graph here.
    while_op = op.outputs[0].op
    cond_graph = _get_graph(while_op, "cond")
    body_graph = _get_graph(while_op, "body")
    orig_num_params = len(body_graph.outputs)
  
    maximum_iterations = op.inputs[1]
    parallel_iterations = op.get_attr("parallel_iterations")
  
    try:
      num_original_outputs = while_op.get_attr("_num_original_outputs")
    except:  # pylint: disable=bare-except
      num_original_outputs = len(while_op.outputs)
  
    num_intermediates = len(while_op.outputs) - num_original_outputs
    grads = [
        _preprocess_grad(grad, body_out, while_out)  # pylint: disable=g-complex-comprehension
        for grad, body_out, while_out in zip(
            grads[:num_original_outputs],
            body_graph.outputs[:num_original_outputs],
            while_op.outputs[:num_original_outputs])
    ] + [None] * num_intermediates
  
    # We compute the gradient for the sub-graph between trainable ys and xs
    # with non-None incoming gradients. We later pad the None's to the list of
    # outputs.
    ys, xs, non_none_grads = zip(*[(y, x, grad) for (y, x, grad) in zip(
        body_graph.outputs, body_graph.inputs, grads) if grad is not None])
  
    body_grad_graph, args = _create_grad_func(
        ys, xs, non_none_grads, cond_graph, body_graph,
        util.unique_grad_fn_name(body_graph.name), op, maximum_iterations)
  
    if body_grad_graph.while_op_needs_rewrite:
      # Modify 'op' to output the intermediate accumulators needed by the grad
      # function.
      # NOTE(skyewm): if there are any active sessions, this modification to `op`
      # may make them unrunnable!
  
      cond_graph.name += "_rewritten"
      body_graph.name += "_rewritten"
  
      new_inputs = body_grad_graph.empty_tensor_lists
      new_outputs = body_graph.outputs[orig_num_params:]
  
      while_op._set_func_attr("cond", util.create_new_tf_function(cond_graph))
      while_op._set_func_attr("body", util.create_new_tf_function(body_graph))
      while_op._set_type_list_attr("T", body_graph.output_types)
      while_op._set_shape_list_attr("output_shapes", body_graph.output_shapes)
      while_op._add_while_inputs(new_inputs)
      while_op._add_outputs([t.dtype for t in new_outputs],
                            [t.shape for t in new_outputs])
      _copy_handle_data(new_outputs, op.outputs[orig_num_params:])
  
    # Do not ingore grads wrt extra outputs when computing higher order
    # derivatives.
    while_op._set_attr("_num_original_outputs",
                       attr_value_pb2.AttrValue(i=len(while_op.outputs)))
  
    captured_inputs = _resolve_grad_captures(body_graph, body_grad_graph,
                                             while_op)
    loop_vars = args + captured_inputs
  
    # This modifies body_grad_graph.
    loop_vars = while_v2_indexed_slices_rewriter.rewrite_grad_indexed_slices(
        grads, body_grad_graph, loop_vars, while_op.inputs)
  
    def grad_cond(counter, unused_maximum_iterations_arg, forward_loop_iters,
                  *unused_args):
      return counter < forward_loop_iters
  
    grad_cond_name = util.unique_grad_fn_name(op.get_attr("cond").name)
    cond_grad_graph = func_graph_module.func_graph_from_py_func(
        grad_cond_name, grad_cond, loop_vars, {},
        func_graph=util.WhileCondFuncGraph(grad_cond_name))
  
    _check_num_inputs_outputs(cond_grad_graph, body_grad_graph, len(loop_vars))
  
    outputs = gen_functional_ops._while(
        loop_vars,
        util.create_new_tf_function(cond_grad_graph),
        util.create_new_tf_function(body_grad_graph),
        output_shapes=[t.shape for t in body_grad_graph.outputs],
        parallel_iterations=parallel_iterations,
        name="%s_grad" % while_op.name)
    grad_op = outputs[0].op
  
    _copy_handle_data(body_grad_graph.outputs, outputs)
    util.maybe_set_lowering_attr(grad_op)
    util.maybe_propagate_compile_time_consts_in_xla(grad_op)
  
    # See comment in while_loop.
    outputs = [array_ops.identity(t) for t in outputs]
    return _get_structured_grad_output(outputs, grads, body_grad_graph)
function_tokens: def _WhileGrad ( op , * grads ) : # pylint: disable=invalid-name # Note that op is not always the same as while_op because the gradient tape, # for eager mode compatibility, forgets information about the proper op. Since # the loop cannot run in eager mode, however, we can safely introspect into # the graph here. while_op = op . outputs [ 0 ] . op cond_graph = _get_graph ( while_op , "cond" ) body_graph = _get_graph ( while_op , "body" ) orig_num_params = len ( body_graph . outputs ) maximum_iterations = op . inputs [ 1 ] parallel_iterations = op . get_attr ( "parallel_iterations" ) try : num_original_outputs = while_op . get_attr ( "_num_original_outputs" ) except : # pylint: disable=bare-except num_original_outputs = len ( while_op . outputs ) num_intermediates = len ( while_op . outputs ) - num_original_outputs grads = [ _preprocess_grad ( grad , body_out , while_out ) # pylint: disable=g-complex-comprehension for grad , body_out , while_out in zip ( grads [ : num_original_outputs ] , body_graph . outputs [ : num_original_outputs ] , while_op . outputs [ : num_original_outputs ] ) ] + [ None ] * num_intermediates # We compute the gradient for the sub-graph between trainable ys and xs # with non-None incoming gradients. We later pad the None's to the list of # outputs. ys , xs , non_none_grads = zip ( * [ ( y , x , grad ) for ( y , x , grad ) in zip ( body_graph . outputs , body_graph . inputs , grads ) if grad is not None ] ) body_grad_graph , args = _create_grad_func ( ys , xs , non_none_grads , cond_graph , body_graph , util . unique_grad_fn_name ( body_graph . name ) , op , maximum_iterations ) if body_grad_graph . while_op_needs_rewrite : # Modify 'op' to output the intermediate accumulators needed by the grad # function. # NOTE(skyewm): if there are any active sessions, this modification to `op` # may make them unrunnable! cond_graph . name += "_rewritten" body_graph . name += "_rewritten" new_inputs = body_grad_graph . empty_tensor_lists new_outputs = body_graph . outputs [ orig_num_params : ] while_op . _set_func_attr ( "cond" , util . create_new_tf_function ( cond_graph ) ) while_op . _set_func_attr ( "body" , util . create_new_tf_function ( body_graph ) ) while_op . _set_type_list_attr ( "T" , body_graph . output_types ) while_op . _set_shape_list_attr ( "output_shapes" , body_graph . output_shapes ) while_op . _add_while_inputs ( new_inputs ) while_op . _add_outputs ( [ t . dtype for t in new_outputs ] , [ t . shape for t in new_outputs ] ) _copy_handle_data ( new_outputs , op . outputs [ orig_num_params : ] ) # Do not ingore grads wrt extra outputs when computing higher order # derivatives. while_op . _set_attr ( "_num_original_outputs" , attr_value_pb2 . AttrValue ( i = len ( while_op . outputs ) ) ) captured_inputs = _resolve_grad_captures ( body_graph , body_grad_graph , while_op ) loop_vars = args + captured_inputs # This modifies body_grad_graph. loop_vars = while_v2_indexed_slices_rewriter . rewrite_grad_indexed_slices ( grads , body_grad_graph , loop_vars , while_op . inputs ) def grad_cond ( counter , unused_maximum_iterations_arg , forward_loop_iters , * unused_args ) : return counter < forward_loop_iters grad_cond_name = util . unique_grad_fn_name ( op . get_attr ( "cond" ) . name ) cond_grad_graph = func_graph_module . func_graph_from_py_func ( grad_cond_name , grad_cond , loop_vars , { } , func_graph = util . WhileCondFuncGraph ( grad_cond_name ) ) _check_num_inputs_outputs ( cond_grad_graph , body_grad_graph , len ( loop_vars ) ) outputs = gen_functional_ops . _while ( loop_vars , util . create_new_tf_function ( cond_grad_graph ) , util . create_new_tf_function ( body_grad_graph ) , output_shapes = [ t . shape for t in body_grad_graph . outputs ] , parallel_iterations = parallel_iterations , name = "%s_grad" % while_op . name ) grad_op = outputs [ 0 ] . op _copy_handle_data ( body_grad_graph . outputs , outputs ) util . maybe_set_lowering_attr ( grad_op ) util . maybe_propagate_compile_time_consts_in_xla ( grad_op ) # See comment in while_loop. outputs = [ array_ops . identity ( t ) for t in outputs ] return _get_structured_grad_output ( outputs , grads , body_grad_graph )