Python3 type annotation does not work with @tf.function + for loop -> tf.while_loop conversion

[jackshi0912]

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): 18.04
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: n/a
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): 2.0 alpha
• Python version: 3.6.7
• Bazel version (if compiling from source): n/a
• GCC/Compiler version (if compiling from source): n/a
• CUDA/cuDNN version: n/a
• GPU model and memory: n/a

Describe the current behavior
If you use python3 type annotation such as x:tf.Tensor = tf.constant(0) (I aliased tf.Tensor for various shape to keep my sanity for reinforcement learning problems) in a @tf.function and the function contains a for loop to be translated to tf.while_loop (that doesn't even have to use the tensor that's annotated), the code will fail as if you did not turn on eager execution.

Describe the expected behavior
Python3 type hinting should not fail the code.

Code to reproduce the issue

@tf.function
def tf_for_tf_break():
    x: tf.Tensor = tf.constant(0)
    for i in tf.range(5):
        x += i
    return x

print(tf_for_tf_break())

Other info / logs
WARNING: Logging before flag parsing goes to stderr.
W0519 21:35:32.992043 140297958307648 tf_logging.py:161] Entity <function tf_for_tf_break at 0x7f99a9edde18> could not be transformed and will be staged without change. Error details can be found in the logs when running with the env variable AUTOGRAPH_VERBOSITY >= 1. Please report this to the AutoGraph team. Cause: AttributeError during conversion: 'NoneType' object has no attribute '_fields'
Traceback (most recent call last):
File "/home/jackshi/.local/lib/python3.6/site-packages/tensorflow/python/autograph/impl/conversion.py", line 393, in function_to_graph
node = node_to_graph(node, context)
File "/home/jackshi/.local/lib/python3.6/site-packages/tensorflow/python/autograph/impl/conversion.py", line 436, in node_to_graph
node = converter.standard_analysis(node, context, is_initial=True)
File "/home/jackshi/.local/lib/python3.6/site-packages/tensorflow/python/autograph/core/converter.py", line 493, in standard_analysis
graphs = cfg.build(node)
File "/home/jackshi/.local/lib/python3.6/site-packages/tensorflow/python/autograph/pyct/cfg.py", line 813, in build
visitor.visit(node)
File "/usr/lib/python3.6/ast.py", line 253, in visit
return visitor(node)
File "/home/jackshi/.local/lib/python3.6/site-packages/tensorflow/python/autograph/pyct/cfg.py", line 672, in visit_FunctionDef
self.visit(stmt)
File "/usr/lib/python3.6/ast.py", line 253, in visit
return visitor(node)
File "/usr/lib/python3.6/ast.py", line 257, in generic_visit
for field, value in iter_fields(node):
File "/usr/lib/python3.6/ast.py", line 171, in iter_fields
for field in node._fields:
AttributeError: 'NoneType' object has no attribute '_fields'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "/home/jackshi/.local/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 369, in converted_call
experimental_partial_types=partial_types)
File "/home/jackshi/.local/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 513, in to_graph
arg_values, arg_types)
File "/home/jackshi/.local/lib/python3.6/site-packages/tensorflow/python/autograph/impl/conversion.py", line 190, in entity_to_graph
node, name, ns = function_to_graph(o, program_ctx, arg_values, arg_types)
File "/home/jackshi/.local/lib/python3.6/site-packages/tensorflow/python/autograph/impl/conversion.py", line 396, in function_to_graph
raise errors.InternalError('conversion', e)
tensorflow.python.autograph.pyct.errors.InternalError: AttributeError during conversion: 'NoneType' object has no attribute '_fields'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "/home/jackshi/MagneticAccelerator/descrete_optimization/tf_scratch.py", line 12, in
print(tf_for_tf_break())
File "/home/jackshi/.local/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py", line 426, in call
self._initialize(args, kwds, add_initializers_to=initializer_map)
File "/home/jackshi/.local/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py", line 370, in _initialize
*args, **kwds))
File "/home/jackshi/.local/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 1313, in _get_concrete_function_internal_garbage_collected
graph_function, _, _ = self._maybe_define_function(args, kwargs)
File "/home/jackshi/.local/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 1580, in _maybe_define_function
graph_function = self._create_graph_function(args, kwargs)
File "/home/jackshi/.local/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 1512, in _create_graph_function
capture_by_value=self._capture_by_value),
File "/home/jackshi/.local/lib/python3.6/site-packages/tensorflow/python/framework/func_graph.py", line 694, in func_graph_from_py_func
func_outputs = python_func(*func_args, **func_kwargs)
File "/home/jackshi/.local/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py", line 317, in wrapped_fn
return weak_wrapped_fn().wrapped(*args, **kwds)
File "/home/jackshi/.local/lib/python3.6/site-packages/tensorflow/python/framework/func_graph.py", line 686, in wrapper
), args, kwargs)
File "/home/jackshi/.local/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 390, in converted_call
return _call_unconverted(f, args, kwargs)
File "/home/jackshi/.local/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 188, in _call_unconverted
return f(*args, **kwargs)
File "/home/jackshi/MagneticAccelerator/descrete_optimization/tf_scratch.py", line 7, in tf_for_tf_break
for i in tf.range(5):
File "/home/jackshi/.local/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 449, in iter
"Tensor objects are only iterable when eager execution is "
TypeError: Tensor objects are only iterable when eager execution is enabled. To iterate over this tensor use tf.map_fn.

[gadagashwini-zz]

I am able to reproduce the issue on colab with TF 2.0alpha . This is the error I got TypeError: Tensor objects are only iterable when eager execution is enabled. To iterate over this tensor use tf.map_fn.

[tensorflow-bot]

Are you satisfied with the resolution of your issue?
Yes
No