Scales gradients for the backwards pass.

  This might be used to, for example, allow one part of a model to learn at a
  lower rate than the rest.

  WARNING: Think carefully about how your optimizer works. If, for example, you
  use rmsprop, the gradient is always rescaled (with some additional epsilon)
  towards unity. This means `scale_gradient` won't have the effect of
  lowering the learning rate.

  If `scale` is `0.0`, this op reduces to `tf.stop_gradient`. If `scale`
  is `1.0`, this op reduces to `tf.identity`.

  Args:
    net: A `tf.Tensor` or in eager mode a callable that produces a `tf.Tensor`.
    scale: The scale factor for the gradient on the backwards pass.
    name: A name for the operation (optional).

  Returns:
    In graph mode returns a `tf.Tensor` with the same type as the input tensor.
    In eager mode returns a callable wrapping `net` whose gradients are scaled.

  Raises:
    ValueError: If `net` dtype is non-float and `scale` is not zero or one.

def scale_gradient(net, scale, name="scale_gradient"):
  """Scales gradients for the backwards pass.

  This might be used to, for example, allow one part of a model to learn at a
  lower rate than the rest.

  WARNING: Think carefully about how your optimizer works. If, for example, you
  use rmsprop, the gradient is always rescaled (with some additional epsilon)
  towards unity. This means `scale_gradient` won't have the effect of
  lowering the learning rate.

  If `scale` is `0.0`, this op reduces to `tf.stop_gradient`. If `scale`
  is `1.0`, this op reduces to `tf.identity`.

  Args:
    net: A `tf.Tensor` or in eager mode a callable that produces a `tf.Tensor`.
    scale: The scale factor for the gradient on the backwards pass.
    name: A name for the operation (optional).

  Returns:
    In graph mode returns a `tf.Tensor` with the same type as the input tensor.
    In eager mode returns a callable wrapping `net` whose gradients are scaled.

  Raises:
    ValueError: If `net` dtype is non-float and `scale` is not zero or one.
  """
  if tf.executing_eagerly():
    if not callable(net):
      raise ValueError(
          "In eager mode `net` must be a callable (similar to how optimizers "
          "must be used when executing eagerly).")
    return tfe.defun(lambda *a, **k: scale_gradient(net(*a, **k), scale, name))

  if scale == 0.0:
    return tf.stop_gradient(net, name=name)
  elif scale == 1.0:
    return tf.identity(net, name=name)
  else:
    if not net.dtype.is_floating:
      raise ValueError("clip_gradient does not support non-float `net` inputs.")

    with tf.name_scope(name, "scale_gradient", values=[net]):
      dtype = net.dtype.base_dtype  # Convert ref dtypes to regular dtypes.
      scale_tensor = tf.convert_to_tensor(scale, dtype=dtype)

      scale_gradient_op = _scale_gradient_op(dtype)
      output = scale_gradient_op(net, scale_tensor)
      output.set_shape(net.get_shape())

    return output