Datasets, iterators and compiled tf.functions

TensorFlow General Discussion
, ,
1

I am hitting some error involving XLA and data.

This is a reproducible short test:

    data = tf.data.Dataset.from_tensor_slices(list(range(10)))
    data = data.repeat()
    iterator = iter(data)

    @tf.function(jit_compile=False)
    def run(it):
        t = 0
        for _ in range(10):
            x = next(it)
            t += x
        return t

    @tf.function(jit_compile=True)
    def run_compiled(it):
        t = 0
        for _ in range(10):
            x = next(it)
            t += x
        return t

    print(run(iterator))
    print(run_compiled(iterator))

Basically the non-compiled version (run) works fine, while the compiled one (run_compiled) is returning the following error:

tensorflow.python.framework.errors_impl.InternalError: Failed copying input tensor from /job:localhost/replica:0/task:0/device:CPU:0 to /job:localhost/replica:0/task:0/device:GPU:0 in order to run __inference_run_compiled_61: No unary variant device copy function found for direction: 1 and Variant type_index: tensorflow::ResourceDeleter [Op:__inference_run_compiled_61]

I am in the following tf version: 2.7.0-dev20210806

Any idea on what is going on?

Thanks!

3

The error’s a little clearer if you run it without the GPU “Detected unsupported operations when trying to compile graph”:

InvalidArgumentError: Detected unsupported operations when trying to compile
 graph __inference_run_compiled_86[_XlaMustCompile=true,config_proto=3175580994766145631,executor_type=11160318154034397263] 
on XLA_CPU_JIT: IteratorGetNext (No registered 'IteratorGetNext' OpKernel for XLA_CPU_JIT devices
compatible with node {{node IteratorGetNext}}){{node IteratorGetNext}}
1 Like