System information

• Have I written custom code: Yes
• OS Platform and Distribution: Windows 10
• TensorFlow installed from binary: 2.1.0

Describe the current behavior
next(iterator, default) is supposed to give the next element in the iterator or the value given as default if the iterator is at the end.
However, when using the above construction in a function with @tf.function, the default value is not returned and an error (tensorflow.python.framework.errors_impl.OutOfRangeError) is produced when trying to call next on an iterator that is at the end.
When running this code in eager mode, the default value is returned as expected.

Describe the expected behavior
In graph mode the default value should be returned when at the end of an iterator.

Standalone code to reproduce the issue

import tensorflow as tf

x = tf.convert_to_tensor([[1], [2], [3]])
ds = tf.data.Dataset.from_tensor_slices(x)
dsi = iter(ds)


@tf.function # remove this to get the expected behaviour
def func():
    for _ in range(4):
        tf.print(next(dsi, -1))


func()

Output (see below for a full stacktrace):

[1]
[2]
[3]
2020-03-27 18:56:09.523946: W tensorflow/core/common_runtime/base_collective_executor.cc:217] BaseCollectiveExecutor::StartAbort Out of range: End of sequence
	 [[{{node IteratorGetNext_3}}]]

Expected output:

[1]
[2]
[3]
-1

Other info / logs
Colab link: https://colab.research.google.com/drive/1PBxoXiE48aC-bo-aY-Bau1Igt4Aj6OFy

stacktrace.txt