tf.map_fn and Defun clash - tensorflow

Question

I'am trying to define my own operation in tensorflow using the #function.Defun decorator, because i need to equip it with a custom gradient. However, in my code i also needs to call the tf.map_fn function, to apply a specific function to a batch of data. For simplicity i am trying to do something like:
def add_func(x):
return x+1
#function.Defun(tf.float32)
def test(a):
return tf.map_fn(add, a)
with tf.Session() as sess:
a = tf.ones(shape=(6,1))
res = sess.run(test(a))
This code runs fine without the #function.Defun decorator, but using the decorator I get the the following error:
InvalidArgumentError: 25 nodes in a cycle
[[Node: test_8028ca0d_2 = test_8028ca0d[_device="/job:localhost/replica:0/task:0/cpu:0"](ones_4)]]
Caused by op 'test_8028ca0d_2', defined at:
File "C:\Users\Nicki\Anaconda3\lib\site-packages\spyder\utils\ipython\start_kernel.py", line 223, in <module>
main()
File "C:\Users\Nicki\Anaconda3\lib\site-packages\spyder\utils\ipython\start_kernel.py", line 219, in main
kernel.start()
File "C:\Users\Nicki\Anaconda3\lib\site-packages\ipykernel\kernelapp.py", line 474, in start
ioloop.IOLoop.instance().start()
File "C:\Users\Nicki\Anaconda3\lib\site-packages\zmq\eventloop\ioloop.py", line 177, in start
super(ZMQIOLoop, self).start()
File "C:\Users\Nicki\Anaconda3\lib\site-packages\tornado\ioloop.py", line 887, in start
handler_func(fd_obj, events)
File "C:\Users\Nicki\Anaconda3\lib\site-packages\tornado\stack_context.py", line 275, in null_wrapper
return fn(*args, **kwargs)
File "C:\Users\Nicki\Anaconda3\lib\site-packages\zmq\eventloop\zmqstream.py", line 440, in _handle_events
self._handle_recv()
File "C:\Users\Nicki\Anaconda3\lib\site-packages\zmq\eventloop\zmqstream.py", line 472, in _handle_recv
self._run_callback(callback, msg)
File "C:\Users\Nicki\Anaconda3\lib\site-packages\zmq\eventloop\zmqstream.py", line 414, in _run_callback
callback(*args, **kwargs)
File "C:\Users\Nicki\Anaconda3\lib\site-packages\tornado\stack_context.py", line 275, in null_wrapper
return fn(*args, **kwargs)
File "C:\Users\Nicki\Anaconda3\lib\site-packages\ipykernel\kernelbase.py", line 276, in dispatcher
return self.dispatch_shell(stream, msg)
File "C:\Users\Nicki\Anaconda3\lib\site-packages\ipykernel\kernelbase.py", line 228, in dispatch_shell
handler(stream, idents, msg)
File "C:\Users\Nicki\Anaconda3\lib\site-packages\ipykernel\kernelbase.py", line 390, in execute_request
user_expressions, allow_stdin)
File "C:\Users\Nicki\Anaconda3\lib\site-packages\ipykernel\ipkernel.py", line 196, in do_execute
res = shell.run_cell(code, store_history=store_history, silent=silent)
File "C:\Users\Nicki\Anaconda3\lib\site-packages\ipykernel\zmqshell.py", line 501, in run_cell
return super(ZMQInteractiveShell, self).run_cell(*args, **kwargs)
File "C:\Users\Nicki\Anaconda3\lib\site-packages\IPython\core\interactiveshell.py", line 2717, in run_cell
interactivity=interactivity, compiler=compiler, result=result)
File "C:\Users\Nicki\Anaconda3\lib\site-packages\IPython\core\interactiveshell.py", line 2827, in run_ast_nodes
if self.run_code(code, result):
File "C:\Users\Nicki\Anaconda3\lib\site-packages\IPython\core\interactiveshell.py", line 2881, in run_code
exec(code_obj, self.user_global_ns, self.user_ns)
File "<ipython-input-19-c5e48d04d428>", line 1, in <module>
runfile('C:/Users/Nicki/.spyder-py3/temp.py', wdir='C:/Users/Nicki/.spyder-py3')
File "C:\Users\Nicki\Anaconda3\lib\site-packages\spyder\utils\site\sitecustomize.py", line 866, in runfile
execfile(filename, namespace)
File "C:\Users\Nicki\Anaconda3\lib\site-packages\spyder\utils\site\sitecustomize.py", line 102, in execfile
exec(compile(f.read(), filename, 'exec'), namespace)
File "C:/Users/Nicki/.spyder-py3/temp.py", line 156, in <module>
res = sess.run(test(a))
File "C:\Users\Nicki\Anaconda3\lib\site-packages\tensorflow\python\framework\function.py", line 618, in __call__
return _call(self._definition.signature, *args, **kwargs)
File "C:\Users\Nicki\Anaconda3\lib\site-packages\tensorflow\python\framework\function.py", line 271, in _call
compute_shapes=False)
File "C:\Users\Nicki\Anaconda3\lib\site-packages\tensorflow\python\framework\ops.py", line 2240, in create_op
original_op=self._default_original_op, op_def=op_def)
File "C:\Users\Nicki\Anaconda3\lib\site-packages\tensorflow\python\framework\ops.py", line 1128, in __init__
self._traceback = _extract_stack()
InvalidArgumentError (see above for traceback): 25 nodes in a cycle
[[Node: test_8028ca0d_2 = test_8028ca0d[_device="/job:localhost/replica:0/task:0/cpu:0"](ones_4)]]

Answer 1

Just to close the loop on this question, the discussion on GitHub revealed that the problem exists in an older version of TensorFlow (0.12 or earlier). Upgrading to TensorFlow 1.0 fixes the problem.