def use_gradient(grad_f):
  """Decorator for easily setting custom gradients for TensorFlow functions.

  * DO NOT use this function if you need to serialize your graph.
  * This function will cause the decorated function to run slower.

  Example:

    def _foo_grad(op, grad): ...

    @use_gradient(_foo_grad)
    def foo(x1, x2, x3): ...

  Args:
    grad_f: function to use as gradient.

  Returns:
    A decorator to apply to the function you wish to override the gradient of.

  """
  grad_f_name = register_to_random_name(grad_f)

  def function_wrapper(f):
    def inner(*inputs):

      # TensorFlow only supports (as of writing) overriding the gradient of
      # individual ops. In order to override the gardient of `f`, we need to
      # somehow make it appear to be an individual TensorFlow op.
      #
      # Our solution is to create a PyFunc that mimics `f`.
      #
      # In particular, we construct a graph for `f` and run it, then use a
      # stateful PyFunc to stash it's results in Python. Then we have another
      # PyFunc mimic it by taking all the same inputs and returning the stashed
      # output.
      #
      # I wish we could do this without PyFunc, but I don't see a way to have
      # it be fully general.

      state = {"out_value": None}

      # First, we need to run `f` and store it's output.

      out = f(*inputs)

      def store_out(out_value):
        """Store the value of out to a python variable."""
        state["out_value"] = out_value

      store_name = "store_" + f.__name__
      store = tf.py_func(store_out, [out], (), stateful=True, name=store_name)

      # Next, we create the mock function, with an overriden gradient.
      # Note that we need to make sure store gets evaluated before the mock
      # runs.

      def mock_f(*inputs):
        """Mimic f by retrieving the stored value of out."""
        return state["out_value"]

      with tf.control_dependencies([store]):
        with gradient_override_map({"PyFunc": grad_f_name}):
          mock_name = "mock_" + f.__name__
          mock_out = tf.py_func(mock_f, inputs, out.dtype, stateful=True,
                                name=mock_name)
          mock_out.set_shape(out.get_shape())

      # Finally, we can return the mock.

      return mock_out
    return inner
  return function_wrapper
def use_gradient ( grad_f ) : grad_f_name = register_to_random_name ( grad_f ) def function_wrapper ( f ) : def inner ( * inputs ) : state = { "out_value" : None } out = f ( * inputs ) def store_out ( out_value ) : state [ "out_value" ] = out_value store_name = "store_" + f . __name__ store = tf . py_func ( store_out , [ out ] , ( ) , stateful = True , name = store_name ) def mock_f ( * inputs ) : return state [ "out_value" ] with tf . control_dependencies ( [ store ] ) : with gradient_override_map ( { "PyFunc" : grad_f_name } ) : mock_name = "mock_" + f . __name__ mock_out = tf . py_func ( mock_f , inputs , out . dtype , stateful = True , name = mock_name ) mock_out . set_shape ( out . get_shape ( ) ) return mock_out return inner return function_wrapper
Decorator for easily setting custom gradients for TensorFlow functions.

  * DO NOT use this function if you need to serialize your graph.
  * This function will cause the decorated function to run slower.

  Example:

    def _foo_grad(op, grad): ...

    @use_gradient(_foo_grad)
    def foo(x1, x2, x3): ...

  Args:
    grad_f: function to use as gradient.

  Returns:
    A decorator to apply to the function you wish to override the gradient of.