tf.while example is not working in eager mode

[arogozhnikov]

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): yes, see below
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 16.04
• TensorFlow installed from (source or binary): pip
• TensorFlow version (use command below): 1.7, release
• Python version: 3.6
• Bazel version (if compiling from source): NA
• GCC/Compiler version (if compiling from source): NA
• CUDA/cuDNN version: 9.0 / 7.0
• GPU model and memory: GTX TITAN, 6GB
• Exact command to reproduce: see below

Describe the problem

tf.while example is not working in eager mode

Source code / logs

Here is the code

import tensorflow as tf
tf.enable_eager_execution() # the only line added to the example

i = tf.constant(0)
c = lambda i: tf.less(i, 10)
b = lambda i: tf.add(i, 1)
r = tf.while_loop(c, b, [i])

result:
TypeError: Cannot iterate over a scalar tensor.

Reason

b = lambda i: tf.add(i, 1)       # original
b = lambda i: (tf.add(i, 1), )   # fixed, working 

• Body code (b) returns a single scalar, which code later tries to unpack.
• Eager code doesn't support auto-wrapping single item to tuple

Possible solutions:

• require body function to always return tuple
• always check if body output is a single variable

[asimshankar]

Thanks for the report, @akshaym is looking into fixing that.

[danmou]

This is fixed on TF 1.15 but broken on TF 2.0.

[akshaym]

I can reproduce this on 2.0. @saxenasaurabh can you please help take a look?

[datdinhquoc]

On Colab, TF 2, api: https://www.tensorflow.org/api_docs/python/tf/while_loop

Example in api doc:

i = tf.constant(0)
c = lambda i: tf.less(i, 10)
b = lambda i: tf.add(i, 1)
r = tf.while_loop(c, b, [i])

TF 2 on Eager mode by default:

TypeError: Cannot iterate over a scalar tensor.

My work-around (wrap in autograph @tf.function):

@tf.function
def f():
  i = tf.constant(0)
  c = lambda i: tf.less(i, 10)
  b = lambda i: tf.add(i, 1)
  r = tf.while_loop(c, b, [i])
  return r

[yaowang1111]

Has this bug been solved? I run tf2.1 and the problem is still there.

[mohantym]

I was able to replicate and resolve this issue in 2.7

[Engineero]

It turns out this has been causing a headache for me through several versions of TF 2.x. I anxiously await a resolution.

[google-ml-butler]

This issue has been automatically marked as stale because it has no recent activity. It will be closed if no further activity occurs. Thank you.

[google-ml-butler]

Closing as stale. Please reopen if you'd like to work on this further.