Using tf.custom_gradient in tensorflow r1.8

Question

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Y
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 16.04
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): r1.8
• Python version: 2.7.14
• GCC/Compiler version (if compiling from source): 5.4
• CUDA/cuDNN version: 8.0/7.0
• GPU model and memory: GTX1080, 8G
• Bazel version: N/A
• Exact command to reproduce: python test_script.py

Describe the problem

Hello, I'm trying to make a custom_gradient op using the function of tf.custom_gradient. I made my test code based on the API explanation online. However, it seems there is a problem in the custom_gradient function. Thanks!

Source code / logs

import tensorflow as tf
import numpy as np

@tf.custom_gradient
def log1pexp(x):
  e = tf.exp(x)
  def grad(dy):
    return dy * (1 - 1 / (1 + e))
  return tf.log(1 + e), grad

x = tf.constant(100.)
f = tf.custom_gradient(log1pexp)

y, dy = f(x)

sess = tf.Session()
print (y.eval(session=sess), y.eval(session=sess).shape)

File "/home/local/home/research/DL/unit_tests/tf_test_custom_grad.py", line 14, in <module>
    y, dy = f(x)
  File "/usr/local/lib/python2.7/dist-packages/tensorflow/python/ops/custom_gradient.py", line 111, in decorated
    return _graph_mode_decorator(f, *args, **kwargs)
  File "/usr/local/lib/python2.7/dist-packages/tensorflow/python/ops/custom_gradient.py", line 132, in _graph_mode_decorator
    result, grad_fn = f(*args)
  File "/usr/local/lib/python2.7/dist-packages/tensorflow/python/framework/ops.py", line 439, in __iter__
"Tensor objects are not iterable when eager execution is not "
TypeError: Tensor objects are not iterable when eager execution is not enabled. To iterate over this tensor use tf.map_fn.

Answer 1

If you just want to test the code in the documentation, here is the way.

The following code will give the instable [nan] result:

import tensorflow as tf

def log1pexp(x):
    return tf.log(1 + tf.exp(x))

x = tf.constant(100.)
y = log1pexp(x)
dy = tf.gradients(y, x)

with tf.Session() as sess:
    print(sess.run(dy))

And the following code will give the correct result [1.0]:

import tensorflow as tf

@tf.custom_gradient
def log1pexp(x):
    e = tf.exp(x)
    def grad(dy):
        return dy * (1 - 1 / (1 + e))
    return tf.log(1 + e), grad

x = tf.constant(100.)
y = log1pexp(x)
dy = tf.gradients(y, x)

with tf.Session() as sess:
    print(sess.run(dy))

Details:

The main problem here is that you are trying to decorate log1pexp twice in your code: once with @tf.custom_gradient and once with f = tf.custom_gradient(log1pexp). In python, @tf.custom_gradient here is equivalent to log1pexp = tf.custom_gradient(log1pexp). You should do this only once, especially here for the following reason.

tf.custom_gradient needs to call the function being pass to it to get both the function output and the gradient, i.e. expecting two returns. During decoration, everything works as expected because log1pexp returns tf.log(1 + e) and grad. After decorating log1pexp, log1pexp (returned by tf.custom_gradient) becomes a new function which returns only one tensor tf.log(1 + e). When you do f = tf.custom_gradient(log1pexp) after decorating log1pexp, tf.custom_gradient can only get one return which is the single tensor tf.log(1 + e). It will try to split this tensor into two by iterating this returned tensor. But it is wrong and is not allowed as the error message stated:

Tensor objects are not iterable when eager execution is not enabled.

You should not decorate log1pexp twice anyway. But this is why you got this error. One more thing to mention, your code will trigger another error for the same reason even if you removed @tf.custom_gradient. After removing @tf.custom_gradient, the line f = tf.custom_gradient(log1pexp) should work as expected. But f is a function returning only one tensor. y, dy = f(x) is wrong and will not work.