parameters: (
      vspace,
      tape,
      target,
      sources,
      output_gradients=None)
return_statement: return result
docstring: Computes gradients from the imperatively defined tape on top of the stack.
  
    Works by filtering the tape, computing how many downstream usages are of each
    tensor and entry, and repeatedly applying backward functions until we have
    gradients for all sources.
  
    Args:
     vspace: the vector space in which to differentiate.
     tape: the gradient tape which stores the trace.
     target: either a Tensor or list of Tensors to be differentiated.
     sources: list of Tensors for which we want gradients
     output_gradients: if not None, a list of gradient provided for each Target,
      or None if we are to use the target's computed downstream gradient.
  
    Returns:
     the gradient wrt each of the sources.
  
    Raises:
      RuntimeError: if something goes wrong.
      ValueError: if there is no sequence of differentiable operations connecting
       a source and any target Tensor. This can happen either if the target is
       not computed based on the source, if the tracing was set up incorrectly,
       or if only non-differentiable functions of the source were used in the
       computation of target.
docstring_summary: Computes gradients from the imperatively defined tape on top of the stack.
docstring_tokens: Computes gradients from the imperatively defined tape on top of the stack .
function: def imperative_grad(
      vspace,
      tape,
      target,
      sources,
      output_gradients=None):
    """Computes gradients from the imperatively defined tape on top of the stack.
  
    Works by filtering the tape, computing how many downstream usages are of each
    tensor and entry, and repeatedly applying backward functions until we have
    gradients for all sources.
  
    Args:
     vspace: the vector space in which to differentiate.
     tape: the gradient tape which stores the trace.
     target: either a Tensor or list of Tensors to be differentiated.
     sources: list of Tensors for which we want gradients
     output_gradients: if not None, a list of gradient provided for each Target,
      or None if we are to use the target's computed downstream gradient.
  
    Returns:
     the gradient wrt each of the sources.
  
    Raises:
      RuntimeError: if something goes wrong.
      ValueError: if there is no sequence of differentiable operations connecting
       a source and any target Tensor. This can happen either if the target is
       not computed based on the source, if the tracing was set up incorrectly,
       or if only non-differentiable functions of the source were used in the
       computation of target.
    """
    tensor_to_op, op_to_entry = tape.export()
    # This overwrites the op_to_entry variable, which will release all memory used
    # to keep traces that are irrelevant to the gradient computation we're doing
    # here.
    id_sources = [vspace.tensor_id(t) for t in sources]
    tensor_usage_counts, op_to_entry, op_missing_tensor = _prepare_backprop(
        vspace, target, tensor_to_op, op_to_entry, id_sources)
    ready_ops = _initialize_backprop_stack(op_to_entry, op_missing_tensor)
    gradients = _initial_gradients(vspace, target, output_gradients,
                                   tensor_usage_counts)
    gradients_size = dict()
    # Now exhaust the backprop stack
    while ready_ops:
      op = ready_ops.pop()
      op_trace = op_to_entry.pop(op)
      out_gradients = [gradients.pop(t, None) for t in op_trace.output_ids]
  
      # Cache the last used zero tensor. We reuse it if the next one
      # we need is of the same shape and dtype. This is very helpful in
      # large splits and should have negligible overhead in other cases.
      last_shape_and_dtype = None
      last_zeros = None
      for i in range(len(out_gradients)):
        if out_gradients[i] is None:
          # TODO(apassos) this should be in the right device
          none_indices = _grad_fn_accepts_none_for_indices.get(
              op_trace.op_type, None)
          if none_indices is None or i not in none_indices:
            shape_and_dtype = op_trace.output_shape_and_dtype[i]
            if shape_and_dtype != last_shape_and_dtype:
              last_shape_and_dtype = shape_and_dtype
              last_zeros = vspace.zeros(*shape_and_dtype)
            out_gradients[i] = last_zeros
        else:
          out_gradients[i] = vspace.aggregate_fn(out_gradients[i])
  
      in_gradients = op_trace.backward_function(*(out_gradients))
      for i, t in enumerate(op_trace.input_ids):
        if in_gradients[i] is not None:
          t_grads = gradients.setdefault(t, [])
          t_grads.append(in_gradients[i])
          if len(t_grads) >= _MIN_AGGREGATE_COUNT:
            if t not in gradients_size:
              gradients_size[t] = vspace.num_elements_fn(t_grads[-1])
            size = gradients_size[t]
  
            if len(t_grads) * size * 4 > _MIN_AGGREGATE_BYTES:
              t_grads[:] = [vspace.aggregate_fn(t_grads)]
        if tensor_usage_counts.get(t, 0) > 0:
          tensor_usage_counts[t] -= 1
          if (t in tensor_to_op
              and tensor_usage_counts[t] == 0
              and t not in id_sources):
            in_op = tensor_to_op[t]
            if in_op is None or in_op == -1:
              continue
            if op_missing_tensor.get(in_op, 0) > 0:
              op_missing_tensor[in_op] -= 1
              if op_missing_tensor.get(in_op, 0) == 0:
                ready_ops.append(in_op)
    result = []
    for i, s in enumerate(sources):
      g = gradients.get(vspace.tensor_id(s), None)
      if g is None:
        result.append(None)
      else:
        result.append(vspace.aggregate_fn(g))
    return result
function_tokens: def imperative_grad ( vspace , tape , target , sources , output_gradients = None ) : tensor_to_op , op_to_entry = tape . export ( ) # This overwrites the op_to_entry variable, which will release all memory used # to keep traces that are irrelevant to the gradient computation we're doing # here. id_sources = [ vspace . tensor_id ( t ) for t in sources ] tensor_usage_counts , op_to_entry , op_missing_tensor = _prepare_backprop ( vspace , target , tensor_to_op , op_to_entry , id_sources ) ready_ops = _initialize_backprop_stack ( op_to_entry , op_missing_tensor ) gradients = _initial_gradients ( vspace , target , output_gradients , tensor_usage_counts ) gradients_size = dict ( ) # Now exhaust the backprop stack while ready_ops : op = ready_ops . pop ( ) op_trace = op_to_entry . pop ( op ) out_gradients = [ gradients . pop ( t , None ) for t in op_trace . output_ids ] # Cache the last used zero tensor. We reuse it if the next one # we need is of the same shape and dtype. This is very helpful in # large splits and should have negligible overhead in other cases. last_shape_and_dtype = None last_zeros = None for i in range ( len ( out_gradients ) ) : if out_gradients [ i ] is None : # TODO(apassos) this should be in the right device none_indices = _grad_fn_accepts_none_for_indices . get ( op_trace . op_type , None ) if none_indices is None or i not in none_indices : shape_and_dtype = op_trace . output_shape_and_dtype [ i ] if shape_and_dtype != last_shape_and_dtype : last_shape_and_dtype = shape_and_dtype last_zeros = vspace . zeros ( * shape_and_dtype ) out_gradients [ i ] = last_zeros else : out_gradients [ i ] = vspace . aggregate_fn ( out_gradients [ i ] ) in_gradients = op_trace . backward_function ( * ( out_gradients ) ) for i , t in enumerate ( op_trace . input_ids ) : if in_gradients [ i ] is not None : t_grads = gradients . setdefault ( t , [ ] ) t_grads . append ( in_gradients [ i ] ) if len ( t_grads ) >= _MIN_AGGREGATE_COUNT : if t not in gradients_size : gradients_size [ t ] = vspace . num_elements_fn ( t_grads [ - 1 ] ) size = gradients_size [ t ] if len ( t_grads ) * size * 4 > _MIN_AGGREGATE_BYTES : t_grads [ : ] = [ vspace . aggregate_fn ( t_grads ) ] if tensor_usage_counts . get ( t , 0 ) > 0 : tensor_usage_counts [ t ] -= 1 if ( t in tensor_to_op and tensor_usage_counts [ t ] == 0 and t not in id_sources ) : in_op = tensor_to_op [ t ] if in_op is None or in_op == - 1 : continue if op_missing_tensor . get ( in_op , 0 ) > 0 : op_missing_tensor [ in_op ] -= 1 if op_missing_tensor . get ( in_op , 0 ) == 0 : ready_ops . append ( in_op ) result = [ ] for i , s in enumerate ( sources ) : g = gradients . get ( vspace . tensor_id ( s ) , None ) if g is None : result . append ( None ) else : result . append ( vspace . aggregate_fn ( g ) ) return result