code: 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
  
imports: from contextlib import contextmanager
  import numpy
  import tensorflow
  import uuid
label_desc: Use image node_id 3 for calling a global function with example usage: use_gradient(grad_f) and returns: function_wrapper, inner, mock_out, state