torch/csrc/jit/autodiff.cpp
// Takes a grad_desc.f returned from `addReverseInline` and splits off the
// reverse_block into its own graph, storing it in df.
// All intermediates needed in the second stage are added to
// outputs of f, and taken as inputs in df. For a more
// detailed description see Note [Gradient graphs] in autodiff.h.
// This function also initializes the fields in grad_desc that were undefined after
// `addReverseInline` (and extends `df_input_vjps` with vjps for captured temporaries).
Takes a grad_desc.f returned from `addReverseInline` and splits off the
reverse_block into its own graph, storing it in df.
All intermediates needed in the second stage are added to
outputs of f, and taken as inputs in df. For a more
detailed description see Note [Gradient graphs] in autodiff.h.
This function also initializes the fields in grad_desc that were undefined after
`addReverseInline` (and extends `df_input_vjps` with vjps for captured temporaries).
Takes a grad_desc . f returned from ` addReverseInline ` and splits off the reverse_block into its own graph storing it in df . All intermediates needed in the second stage are added to outputs of f and taken as inputs in df . For a more detailed description see Note [ Gradient graphs ] in autodiff . h . This function also initializes the fields in grad_desc that were undefined after ` addReverseInline ` ( and extends ` df_input_vjps ` with vjps for captured temporaries ) .
static void lambdaLiftReverse(Gradient& grad_desc, ReverseDetails& rev_info) {
  auto & graph = *grad_desc.f;
  auto primal_block = graph.block();
  auto reverse_block = rev_info.reverse_block;
  value_set reverse_captures_set;
  value_list reverse_captures; 
  auto check_uses = [&](Value *v) {
    for (auto use : v->uses()) {
      if (use.user->owningBlock() == primal_block)
        continue;
      if ( reverse_captures_set.emplace(v).second) {
        reverse_captures.push_back(v);
      }
    }
  };
  for (Value * input : graph.inputs()) {
    if (input->stage() != 0) break;
    check_uses(input);
  }
  for (Node * node : graph.nodes()) {
    if (node->stage() != 0) break;
    for (Value * output : node->outputs())
      check_uses(output);
  }
   Primal outputs:
     [original outputs], [temporaries]
   Reverse inputs:
     [output vjps (aka grad_outputs)], [temporary vjps]
     [captured primal values, in topological order],
   -- Construct primal_outputs, df_input_captures, f_real_outputs ----
  grad_desc.f_real_outputs = graph.outputs().size();
  std::unordered_map<Value*, size_t> orig_primal_outputs_idx;
  std::unordered_map<Value*, size_t> orig_primal_inputs_idx;
   NOTE: we use emplace to avoid replacing an existing index if an output is repeated
  for (size_t i = 0, num_outputs = graph.outputs().size(); i < num_outputs; ++i)
    orig_primal_outputs_idx.emplace(graph.outputs()[i], i);
  for (size_t i = 0, num_inputs = graph.inputs().size(); i < num_inputs; ++i)
    orig_primal_inputs_idx[graph.inputs()[i]] = i;
   NB: reverse_captures are already deduplicated, and in topo order
  for (Value * capture_val : reverse_captures) {
     If it's already an output we don't have to add anything,
     but register the fact that it needs to be captured.
    if (orig_primal_outputs_idx.count(capture_val) > 0) {
      grad_desc.df_input_captured_outputs.push_back(orig_primal_outputs_idx[capture_val]);
     If it's an input, we could add it as an output but in fact it's
     more efficient to use a special kind of capture.
    } else if (orig_primal_inputs_idx.count(capture_val) > 0) {
      grad_desc.df_input_captured_inputs.push_back(orig_primal_inputs_idx.at(capture_val));
     Otherwise it's just a regular intermediate value that we need to add as an output
    } else {
       we need to create a new temporary output for this capture because it wasn't availiable.
      graph.registerOutput(capture_val);
      grad_desc.df_input_captured_outputs.emplace_back(graph.outputs().size() - 1);
    }
  }
   -- Add VJPs for temporaries, adjust df_input_vjps -------------------------
   NB [possible optimization]: use the newly added vjp input as soon as the first
   vjp for that value is generated, to reduce the lifespan of this input
   (currently we add it to the final vjp after all adds).
  for (size_t i = grad_desc.f_real_outputs; i < graph.outputs().size(); ++i) {
    Value * tmp = graph.outputs().at(i);
     Add VJP inputs only for intermediates that actually required grad.
    if (rev_info.requires_grad_set.count(tmp) == 0) continue;
    Value * tmp_vjp_in = reverse_block->addInput()->setType(tmp->type());
    Value * tmp_vjp_prev = rev_info.grad_map.at(tmp);
     This is quite weird because we can't first make a sum and then replace all uses
     of tmp_vjp_prev (that would replace its use in the sum too!), so we create an
     incorrect sum that doesn't use prev vjp, replace uses, and fix the sum.
    Value * new_vjp = createAutogradAdd(tmp_vjp_in, tmp_vjp_in);
    new_vjp->node()->moveAfter(tmp_vjp_prev->node());
    tmp_vjp_prev->replaceAllUsesWith(new_vjp);
    new_vjp->node()->replaceInput(1, tmp_vjp_prev);
    grad_desc.df_input_vjps.emplace_back(i);
  }
   add the captures as formal arguments to the reverse_block
   afterward inputs: [output vjps][temporary vjps][captures]
   construct a map from captured 'value' to the index in the input list
   used to extract this block into its own function
  std::unordered_map<Value*, size_t> capture_to_formal_index;
  const auto & add_capture = [&](Value * captured) {
    capture_to_formal_index[captured] = reverse_block->inputs().size();
    reverse_block->addInput()->copyMetadata(captured);
  };
  for(auto & offset : grad_desc.df_input_captured_inputs)
    add_capture(graph.inputs()[offset]);
  for(auto & offset : grad_desc.df_input_captured_outputs)
    add_capture(graph.outputs()[offset]);
  grad_desc.df = std::make_shared<Graph>();
  grad_desc.df->block()->cloneFrom(reverse_block, [&](Value* v) {
    return grad_desc.df->inputs()[capture_to_formal_index.at(v)];
  });
   reverse_node was just to hold onto reverse_block in a debuggable way
   we can remove it now.
  reverse_block->owningNode()->destroy();
}
static void lambdaLiftReverse ( Gradient & grad_desc , ReverseDetails & rev_info ) { auto & graph = * grad_desc . f ; auto primal_block = graph . block ( ) ; auto reverse_block = rev_info . reverse_block ; value_set reverse_captures_set ; value_list reverse_captures ; auto check_uses = [ & ] ( Value * v ) { for ( auto use : v -> uses ( ) ) { if ( use . user -> owningBlock ( ) == primal_block ) continue ; if ( reverse_captures_set . emplace ( v ) . second ) { reverse_captures . push_back ( v ) ; } } } ; for ( Value * input : graph . inputs ( ) ) { if ( input -> stage ( ) != 0 ) break ; check_uses ( input ) ; } for ( Node * node : graph . nodes ( ) ) { if ( node -> stage ( ) != 0 ) break ; for ( Value * output : node -> outputs ( ) ) check_uses ( output ) ; } grad_desc . f_real_outputs = graph . outputs ( ) . size ( ) ; std :: unordered_map < Value * , size_t > orig_primal_outputs_idx ; std :: unordered_map < Value * , size_t > orig_primal_inputs_idx ; for ( size_t i = 0 , num_outputs = graph . outputs ( ) . size ( ) ; i < num_outputs ; ++ i ) orig_primal_outputs_idx . emplace ( graph . outputs ( ) [ i ] , i ) ; for ( size_t i = 0 , num_inputs = graph . inputs ( ) . size ( ) ; i < num_inputs ; ++ i ) orig_primal_inputs_idx [ graph . inputs ( ) [ i ] ] = i ; for ( Value * capture_val : reverse_captures ) { if ( orig_primal_outputs_idx . count ( capture_val ) > 0 ) { grad_desc . df_input_captured_outputs . push_back ( orig_primal_outputs_idx [ capture_val ] ) ; } else if ( orig_primal_inputs_idx . count ( capture_val ) > 0 ) { grad_desc . df_input_captured_inputs . push_back ( orig_primal_inputs_idx . at ( capture_val ) ) ; } else { graph . registerOutput ( capture_val ) ; grad_desc . df_input_captured_outputs . emplace_back ( graph . outputs ( ) . size ( ) - 1 ) ; } } for ( size_t i = grad_desc . f_real_outputs ; i < graph . outputs ( ) . size ( ) ; ++ i ) { Value * tmp = graph . outputs ( ) . at ( i ) ; if ( rev_info . requires_grad_set . count ( tmp ) == 0 ) continue ; Value * tmp_vjp_in = reverse_block -> addInput ( ) -> setType ( tmp -> type ( ) ) ; Value * tmp_vjp_prev = rev_info . grad_map . at ( tmp ) ; Value * new_vjp = createAutogradAdd ( tmp_vjp_in , tmp_vjp_in ) ; new_vjp -> node ( ) -> moveAfter ( tmp_vjp_prev -> node ( ) ) ; tmp_vjp_prev -> replaceAllUsesWith ( new_vjp ) ; new_vjp -> node ( ) -> replaceInput ( 1 , tmp_vjp_prev ) ; grad_desc . df_input_vjps . emplace_back ( i ) ; } std :: unordered_map < Value * , size_t > capture_to_formal_index ; const auto & add_capture = [ & ] ( Value * captured ) { capture_to_formal_index [ captured ] = reverse_block -> inputs ( ) . size ( ) ; reverse_block -> addInput ( ) -> copyMetadata ( captured ) ; } ; for ( auto & offset : grad_desc . df_input_captured_inputs ) add_capture ( graph . inputs ( ) [ offset ] ) ; for ( auto & offset : grad_desc . df_input_captured_outputs ) add_capture ( graph . outputs ( ) [ offset ] ) ; grad_desc . df = std :: make_shared < Graph > ( ) ; grad_desc . df -> block ( ) -> cloneFrom ( reverse_block , [ & ] ( Value * v ) { return grad_desc . df -> inputs ( ) [ capture_to_formal_index . at ( v ) ] ; } ) ; reverse_block -> owningNode ( ) -> destroy ( ) ; }
Takes a grad_desc.f returned from `addReverseInline` and splits off the
reverse_block into its own graph, storing it in df.
Takes a grad_desc . f returned from ` addReverseInline ` and splits off the reverse_block into its own graph storing it in df .
// -------------------------------------------------------------------------- // 1. Find values of f that need to be captured. // -------------------------------------------------------------------------- // First, we need to find all values that are produced in f, // and used in df. They will need to be added as inputs of the df // and some of them may also need to be appended as outputs of f if // they are not already an input or an output of f // Invariant: topo sorted /* bool unseen = */ // -------------------------------------------------------------------------- // 2. Prepare input/outputs lists for f and df // -------------------------------------------------------------------------- // It's simple to construct primal_inputs/reverse_outputs, // but primal_outputs/reverse_inputs are much more subtle. // Here's a summary of how they are supposed to look like: // // Primal outputs: //   [original outputs], [temporaries] // // Reverse inputs: //   [output vjps (aka grad_outputs)], [temporary vjps] //   [captured primal values, in topological order], // -- Construct primal_outputs, df_input_captures, f_real_outputs ---- // NOTE: we use emplace to avoid replacing an existing index if an output is repeated // NB: reverse_captures are already deduplicated, and in topo order // If it's already an output we don't have to add anything, // but register the fact that it needs to be captured. // If it's an input, we could add it as an output but in fact it's // more efficient to use a special kind of capture. // Otherwise it's just a regular intermediate value that we need to add as an output // we need to create a new temporary output for this capture because it wasn't availiable. // -- Add VJPs for temporaries, adjust df_input_vjps ------------------------- // NB [possible optimization]: use the newly added vjp input as soon as the first // vjp for that value is generated, to reduce the lifespan of this input // (currently we add it to the final vjp after all adds). // Add VJP inputs only for intermediates that actually required grad. // This is quite weird because we can't first make a sum and then replace all uses // of tmp_vjp_prev (that would replace its use in the sum too!), so we create an // incorrect sum that doesn't use prev vjp, replace uses, and fix the sum. // add the captures as formal arguments to the reverse_block // afterward inputs: [output vjps][temporary vjps][captures] // construct a map from captured 'value' to the index in the input list // used to extract this block into its own function // reverse_node was just to hold onto reverse_block in a debuggable way // we can remove it now.