tensorflow: [autograph List comprehension issue] OperatorNotAllowedInGraphError error in Tensorflow 2.0

System information

• OS Platform and Distribution Mac os (10.14.6)
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): 2.0.0rc1
• Python version: Python 3.6.4

I’m learn tensorflow2.0 from official tutorials.I can understand the result from below code.

def square_if_positive(x):
  return [i ** 2 if i > 0 else i for i in x]
square_if_positive(range(-5, 5))

# result
[-5, -4, -3, -2, -1, 0, 1, 4, 9, 16]

But if I change the inputs with tensor not python code, just like this, I got a error

@tf.function
def square_if_positive(x):
  return [i ** 2 if i > 0 else i for i in x]
square_if_positive(tf.range(-5, 5))

~/tf2_workspace/tf2.0/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py in _create_graph_function(self, args, kwargs, override_flat_arg_shapes)
   1984             arg_names=arg_names,
   1985             override_flat_arg_shapes=override_flat_arg_shapes,
-> 1986             capture_by_value=self._capture_by_value),
   1987         self._function_attributes,
   1988         # Tell the ConcreteFunction to clean up its graph once it goes out of

~/tf2_workspace/tf2.0/lib/python3.6/site-packages/tensorflow_core/python/framework/func_graph.py in func_graph_from_py_func(name, python_func, args, kwargs, signature, func_graph, autograph, autograph_options, add_control_dependencies, arg_names, op_return_value, collections, capture_by_value, override_flat_arg_shapes)
    851                                           converted_func)
    852 
--> 853       func_outputs = python_func(*func_args, **func_kwargs)
    854 
    855       # invariant: `func_outputs` contains only Tensors, CompositeTensors,

~/tf2_workspace/tf2.0/lib/python3.6/site-packages/tensorflow_core/python/eager/def_function.py in wrapped_fn(*args, **kwds)
    323         # __wrapped__ allows AutoGraph to swap in a converted function. We give
    324         # the function a weak reference to itself to avoid a reference cycle.
--> 325         return weak_wrapped_fn().__wrapped__(*args, **kwds)
    326     weak_wrapped_fn = weakref.ref(wrapped_fn)
    327 

~/tf2_workspace/tf2.0/lib/python3.6/site-packages/tensorflow_core/python/framework/func_graph.py in wrapper(*args, **kwargs)
    841           except Exception as e:  # pylint:disable=broad-except
    842             if hasattr(e, "ag_error_metadata"):
--> 843               raise e.ag_error_metadata.to_exception(type(e))
    844             else:
    845               raise

OperatorNotAllowedInGraphError: in converted code:

    <ipython-input-37-6c17f29a3443>:3 square_if_positive  *
        return [i**2 if i > 0 else i for i in x]
    /Users/zhangpan/tf2_workspace/tf2.0/lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py:547 __iter__
        self._disallow_iteration()
    /Users/zhangpan/tf2_workspace/tf2.0/lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py:540 _disallow_iteration
        self._disallow_when_autograph_enabled("iterating over `tf.Tensor`")
    /Users/zhangpan/tf2_workspace/tf2.0/lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py:518 _disallow_when_autograph_enabled
        " decorating it directly with @tf.function.".format(task))

    OperatorNotAllowedInGraphError: iterating over `tf.Tensor` is not allowed: AutoGraph did not convert this function. Try decorating it directly with @tf.function.

I can’t find any specifications about this error. I think the real reason is not “iterating over tf.Tensor is not allowed” . Becase I can write like this.

@tf.function
def square_if_positive(x):
    for i in x:
        if i>0:
            tf.print(i**2)
        else:
            tf.print(i)

square_if_positive(tf.range(10))

iterate over tensor just like above code. So my question is what’s the real reason about this error? Any suggestions will help me. I really can’t understand this error through I read a lot of materials.