def _recompute_grad(fn, args):
  """See recompute_grad."""

  cached_vs = []
  cached_arg_scope = []

  def grad_fn(inputs, variables, outputs, output_grads):
    """Recompute outputs for gradient computation."""
    del outputs
    variables = [underlying_variable_ref(v) for v in variables]
    # Recompute outputs
    with tf.control_dependencies(output_grads):
      with tf.contrib.framework.arg_scope(cached_arg_scope[0]):
        with tf.variable_scope(cached_vs[0], reuse=True):
          outputs = fn(*inputs)

    if not isinstance(outputs, (list, tuple)):
      outputs = [outputs]
    outputs = list(outputs)
    grads = tf.gradients(outputs, inputs + variables, output_grads)
    grad_inputs = grads[:len(inputs)]
    grad_vars = grads[len(inputs):]
    # TODO(rsepassi): Make fn_with_custom_grad work with bfloat16.
    # If the input gradients are bfloat16, it's assumed the variables are
    # bfloat16. This is a hack to ensure that grad_vars are the right type.
    if grad_inputs[0].dtype == tf.bfloat16:
      grad_vars = [tf.cast(grad_var, tf.bfloat16) for grad_var in grad_vars]
    return grad_inputs, grad_vars

  @fn_with_custom_grad(grad_fn)
  def fn_with_recompute(*args):
    cached_vs.append(tf.get_variable_scope())
    cached_arg_scope.append(tf.contrib.framework.current_arg_scope())
    return fn(*args)

  return fn_with_recompute(*args)
def _recompute_grad ( fn , args ) : cached_vs = [ ] cached_arg_scope = [ ] def grad_fn ( inputs , variables , outputs , output_grads ) : del outputs variables = [ underlying_variable_ref ( v ) for v in variables ] with tf . control_dependencies ( output_grads ) : with tf . contrib . framework . arg_scope ( cached_arg_scope [ 0 ] ) : with tf . variable_scope ( cached_vs [ 0 ] , reuse = True ) : outputs = fn ( * inputs ) if not isinstance ( outputs , ( list , tuple ) ) : outputs = [ outputs ] outputs = list ( outputs ) grads = tf . gradients ( outputs , inputs + variables , output_grads ) grad_inputs = grads [ : len ( inputs ) ] grad_vars = grads [ len ( inputs ) : ] if grad_inputs [ 0 ] . dtype == tf . bfloat16 : grad_vars = [ tf . cast ( grad_var , tf . bfloat16 ) for grad_var in grad_vars ] return grad_inputs , grad_vars @ fn_with_custom_grad ( grad_fn ) def fn_with_recompute ( * args ) : cached_vs . append ( tf . get_variable_scope ( ) ) cached_arg_scope . append ( tf . contrib . framework . current_arg_scope ( ) ) return fn ( * args ) return fn_with_recompute ( * args )
See recompute_grad.