def _scale_gradient_op(dtype):
  """Create an op that scales gradients using a Defun.

  The tensorflow Defun decorator creates an op and tensorflow caches these ops
  automatically according to `func_name`. Using a Defun decorator twice with the
  same `func_name` does not create a new op, instead the cached op is used.

  This method produces a new op the first time it is called with a given `dtype`
  argument, and then uses the cached op each time it is called after that with
  the same `dtype`. The scale value is given as an argument for the forward pass
  method so that it can be used in the backwards pass.

  Args:
    dtype: the dtype of the net whose gradient is being scaled.

  Returns:
    The op that scales gradients.
  """

  def scale_gradient_backward(op, grad):
    scale = op.inputs[1]
    scaled_grad = grad * scale
    return scaled_grad, None

  # Note that if the forward pass implementation involved the creation of ops,
  # _scale_gradient_op would require some memoization mechanism.
  def scale_gradient_forward(x, scale):
    del scale  # Unused.
    return x

  func_name = "ScaleGradient_{}".format(dtype.name)
  return function.Defun(
      dtype, dtype,
      python_grad_func=scale_gradient_backward,
      func_name=func_name)(scale_gradient_forward)
def _scale_gradient_op ( dtype ) : def scale_gradient_backward ( op , grad ) : scale = op . inputs [ 1 ] scaled_grad = grad * scale return scaled_grad , None def scale_gradient_forward ( x , scale ) : del scale return x func_name = "ScaleGradient_{}" . format ( dtype . name ) return function . Defun ( dtype , dtype , python_grad_func = scale_gradient_backward , func_name = func_name ) ( scale_gradient_forward )
Create an op that scales gradients using a Defun.

  The tensorflow Defun decorator creates an op and tensorflow caches these ops
  automatically according to `func_name`. Using a Defun decorator twice with the
  same `func_name` does not create a new op, instead the cached op is used.

  This method produces a new op the first time it is called with a given `dtype`
  argument, and then uses the cached op each time it is called after that with
  the same `dtype`. The scale value is given as an argument for the forward pass
  method so that it can be used in the backwards pass.

  Args:
    dtype: the dtype of the net whose gradient is being scaled.

  Returns:
    The op that scales gradients.