tensorflow: Parallel loading of nested ExtensionTypes produces shape mismatch

Click to expand!

Issue Type

Bug

Source

binary

Tensorflow Version

2.10

Custom Code

No

OS Platform and Distribution

Linux Ubuntu 20.04

Mobile device

No response

Python version

3.8

Bazel version

No response

GCC/Compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

Nvidia Geforce 1080 Ti

Current Behaviour?

For some reason tensor `b2` has the shape of `a` and that does not match the definition. This happens not every time I run the exact same code but only sometimes. It seems to work with Tensorflow 2.9 though...

Standalone code to reproduce the issue

import tensorflow as tf


class A(tf.experimental.ExtensionType):
    a1: tf.Tensor

    class Spec:
        def __init__(self):
            self.a1 = tf.TensorSpec(shape=[2, 1], dtype=tf.float16)


class B(tf.experimental.ExtensionType):
    b1: tf.Tensor
    b2: tf.Tensor
    a: A

    class Spec:
        def __init__(self):
            self.b1 = tf.TensorSpec(shape=[2, 5], dtype=tf.float16)
            self.b2 = tf.TensorSpec(shape=[2, 6], dtype=tf.float32)
            self.a = A.Spec()


def loading_fn():
    objects = B(
        tf.zeros(shape=[2, 5], dtype=tf.float16),
        tf.zeros(shape=[2, 6], dtype=tf.float32),
        A(tf.zeros(shape=[2, 1], dtype=tf.float16)),
    )

    return objects


def func(i):
    return tf.py_function(loading_fn, inp=[], Tout=B.Spec())


dataset = tf.data.Dataset.from_tensor_slices([i for i in range(10)])
dataset = dataset.map(func, num_parallel_calls=10)

for _ in range(10):
    for path in iter(dataset):
        pass

Relevant log output

2022-11-10 14:44:32.830413: W tensorflow/core/framework/op_kernel.cc:1768] INVALID_ARGUMENT: ValueError: py_function: func=<function loading_fn at 0x7fc368151a60> returned [B(b1=<tf.Tensor: shape=(2, 5), dtype=float16, numpy=
array([[0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.]], dtype=float16)>, b2=<tf.Tensor: shape=(2, 6), dtype=float32, numpy=
array([[0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0.]], dtype=float32)>, a=A(a1=<tf.Tensor: shape=(2, 6), dtype=float16, numpy=
array([[0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0.]], dtype=float16)>))], which did not match Tout=[B.Spec(b1=TensorSpec(shape=(2, 5), dtype=tf.float16, name=None), b2=TensorSpec(shape=(2, 6), dtype=tf.float32, name=None), a=A.Spec(a1=TensorSpec(shape=(2, 1), dtype=tf.float16, name=None)))].
In particular, B(b1=<tf.Tensor: shape=(2, 5), dtype=float16, numpy=
array([[0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.]], dtype=float16)>, b2=<tf.Tensor: shape=(2, 6), dtype=float32, numpy=
array([[0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0.]], dtype=float32)>, a=A(a1=<tf.Tensor: shape=(2, 6), dtype=float16, numpy=
array([[0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0.]], dtype=float16)>)) is not compatible with B.Spec(b1=TensorSpec(shape=(2, 5), dtype=tf.float16, name=None), b2=TensorSpec(shape=(2, 6), dtype=tf.float32, name=None), a=A.Spec(a1=TensorSpec(shape=(2, 1), dtype=tf.float16, name=None))).
Traceback (most recent call last):

  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/ops/script_ops.py", line 269, in __call__
    return func(device, token, args)

  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/ops/script_ops.py", line 147, in __call__
    outputs = self._call(device, args)

  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/ops/script_ops.py", line 154, in _call
    ret = self._func(*args)

  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/ops/script_ops.py", line 824, in wrapped_func
    raise ValueError(

ValueError: py_function: func=<function loading_fn at 0x7fc368151a60> returned [B(b1=<tf.Tensor: shape=(2, 5), dtype=float16, numpy=
array([[0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.]], dtype=float16)>, b2=<tf.Tensor: shape=(2, 6), dtype=float32, numpy=
array([[0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0.]], dtype=float32)>, a=A(a1=<tf.Tensor: shape=(2, 6), dtype=float16, numpy=
array([[0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0.]], dtype=float16)>))], which did not match Tout=[B.Spec(b1=TensorSpec(shape=(2, 5), dtype=tf.float16, name=None), b2=TensorSpec(shape=(2, 6), dtype=tf.float32, name=None), a=A.Spec(a1=TensorSpec(shape=(2, 1), dtype=tf.float16, name=None)))].
In particular, B(b1=<tf.Tensor: shape=(2, 5), dtype=float16, numpy=
array([[0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.]], dtype=float16)>, b2=<tf.Tensor: shape=(2, 6), dtype=float32, numpy=
array([[0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0.]], dtype=float32)>, a=A(a1=<tf.Tensor: shape=(2, 6), dtype=float16, numpy=
array([[0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0.]], dtype=float16)>)) is not compatible with B.Spec(b1=TensorSpec(shape=(2, 5), dtype=tf.float16, name=None), b2=TensorSpec(shape=(2, 6), dtype=tf.float32, name=None), a=A.Spec(a1=TensorSpec(shape=(2, 1), dtype=tf.float16, name=None))).


2022-11-10 14:44:32.838758: W tensorflow/core/framework/op_kernel.cc:1768] INVALID_ARGUMENT: ValueError: py_function: func=<function loading_fn at 0x7fc368151a60> returned [B(b1=<tf.Tensor: shape=(2, 5), dtype=float16, numpy=
array([[0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.]], dtype=float16)>, b2=<tf.Tensor: shape=(2, 1), dtype=float32, numpy=
array([[0.],
       [0.]], dtype=float32)>, a=A(a1=<tf.Tensor: shape=(2, 1), dtype=float16, numpy=
array([[0.],
       [0.]], dtype=float16)>))], which did not match Tout=[B.Spec(b1=TensorSpec(shape=(2, 5), dtype=tf.float16, name=None), b2=TensorSpec(shape=(2, 6), dtype=tf.float32, name=None), a=A.Spec(a1=TensorSpec(shape=(2, 1), dtype=tf.float16, name=None)))].
In particular, B(b1=<tf.Tensor: shape=(2, 5), dtype=float16, numpy=
array([[0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.]], dtype=float16)>, b2=<tf.Tensor: shape=(2, 1), dtype=float32, numpy=
array([[0.],
       [0.]], dtype=float32)>, a=A(a1=<tf.Tensor: shape=(2, 1), dtype=float16, numpy=
array([[0.],
       [0.]], dtype=float16)>)) is not compatible with B.Spec(b1=TensorSpec(shape=(2, 5), dtype=tf.float16, name=None), b2=TensorSpec(shape=(2, 6), dtype=tf.float32, name=None), a=A.Spec(a1=TensorSpec(shape=(2, 1), dtype=tf.float16, name=None))).
Traceback (most recent call last):

  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/ops/script_ops.py", line 269, in __call__
    return func(device, token, args)

  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/ops/script_ops.py", line 147, in __call__
    outputs = self._call(device, args)

  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/ops/script_ops.py", line 154, in _call
    ret = self._func(*args)

  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/ops/script_ops.py", line 824, in wrapped_func
    raise ValueError(

ValueError: py_function: func=<function loading_fn at 0x7fc368151a60> returned [B(b1=<tf.Tensor: shape=(2, 5), dtype=float16, numpy=
array([[0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.]], dtype=float16)>, b2=<tf.Tensor: shape=(2, 1), dtype=float32, numpy=
array([[0.],
       [0.]], dtype=float32)>, a=A(a1=<tf.Tensor: shape=(2, 1), dtype=float16, numpy=
array([[0.],
       [0.]], dtype=float16)>))], which did not match Tout=[B.Spec(b1=TensorSpec(shape=(2, 5), dtype=tf.float16, name=None), b2=TensorSpec(shape=(2, 6), dtype=tf.float32, name=None), a=A.Spec(a1=TensorSpec(shape=(2, 1), dtype=tf.float16, name=None)))].
In particular, B(b1=<tf.Tensor: shape=(2, 5), dtype=float16, numpy=
array([[0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.]], dtype=float16)>, b2=<tf.Tensor: shape=(2, 1), dtype=float32, numpy=
array([[0.],
       [0.]], dtype=float32)>, a=A(a1=<tf.Tensor: shape=(2, 1), dtype=float16, numpy=
array([[0.],
       [0.]], dtype=float16)>)) is not compatible with B.Spec(b1=TensorSpec(shape=(2, 5), dtype=tf.float16, name=None), b2=TensorSpec(shape=(2, 6), dtype=tf.float32, name=None), a=A.Spec(a1=TensorSpec(shape=(2, 1), dtype=tf.float16, name=None))).


Traceback (most recent call last):
  File "test.py", line 42, in <module>
    for path in iter(dataset):
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/data/ops/iterator_ops.py", line 766, in __next__
    return self._next_internal()
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/data/ops/iterator_ops.py", line 749, in _next_internal
    ret = gen_dataset_ops.iterator_get_next(
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/ops/gen_dataset_ops.py", line 3017, in iterator_get_next
    _ops.raise_from_not_ok_status(e, name)
  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/framework/ops.py", line 7209, in raise_from_not_ok_status
    raise core._status_to_exception(e) from None  # pylint: disable=protected-access
tensorflow.python.framework.errors_impl.InvalidArgumentError: {{function_node __wrapped__IteratorGetNext_output_types_3_device_/job:localhost/replica:0/task:0/device:CPU:0}} ValueError: py_function: func=<function loading_fn at 0x7fc368151a60> returned [B(b1=<tf.Tensor: shape=(2, 5), dtype=float16, numpy=
array([[0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.]], dtype=float16)>, b2=<tf.Tensor: shape=(2, 1), dtype=float32, numpy=
array([[0.],
       [0.]], dtype=float32)>, a=A(a1=<tf.Tensor: shape=(2, 1), dtype=float16, numpy=
array([[0.],
       [0.]], dtype=float16)>))], which did not match Tout=[B.Spec(b1=TensorSpec(shape=(2, 5), dtype=tf.float16, name=None), b2=TensorSpec(shape=(2, 6), dtype=tf.float32, name=None), a=A.Spec(a1=TensorSpec(shape=(2, 1), dtype=tf.float16, name=None)))].
In particular, B(b1=<tf.Tensor: shape=(2, 5), dtype=float16, numpy=
array([[0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.]], dtype=float16)>, b2=<tf.Tensor: shape=(2, 1), dtype=float32, numpy=
array([[0.],
       [0.]], dtype=float32)>, a=A(a1=<tf.Tensor: shape=(2, 1), dtype=float16, numpy=
array([[0.],
       [0.]], dtype=float16)>)) is not compatible with B.Spec(b1=TensorSpec(shape=(2, 5), dtype=tf.float16, name=None), b2=TensorSpec(shape=(2, 6), dtype=tf.float32, name=None), a=A.Spec(a1=TensorSpec(shape=(2, 1), dtype=tf.float16, name=None))).
Traceback (most recent call last):

  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/ops/script_ops.py", line 269, in __call__
    return func(device, token, args)

  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/ops/script_ops.py", line 147, in __call__
    outputs = self._call(device, args)

  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/ops/script_ops.py", line 154, in _call
    ret = self._func(*args)

  File "/usr/local/lib/python3.8/dist-packages/tensorflow/python/ops/script_ops.py", line 824, in wrapped_func
    raise ValueError(

ValueError: py_function: func=<function loading_fn at 0x7fc368151a60> returned [B(b1=<tf.Tensor: shape=(2, 5), dtype=float16, numpy=
array([[0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.]], dtype=float16)>, b2=<tf.Tensor: shape=(2, 1), dtype=float32, numpy=
array([[0.],
       [0.]], dtype=float32)>, a=A(a1=<tf.Tensor: shape=(2, 1), dtype=float16, numpy=
array([[0.],
       [0.]], dtype=float16)>))], which did not match Tout=[B.Spec(b1=TensorSpec(shape=(2, 5), dtype=tf.float16, name=None), b2=TensorSpec(shape=(2, 6), dtype=tf.float32, name=None), a=A.Spec(a1=TensorSpec(shape=(2, 1), dtype=tf.float16, name=None)))].
In particular, B(b1=<tf.Tensor: shape=(2, 5), dtype=float16, numpy=
array([[0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.]], dtype=float16)>, b2=<tf.Tensor: shape=(2, 1), dtype=float32, numpy=
array([[0.],
       [0.]], dtype=float32)>, a=A(a1=<tf.Tensor: shape=(2, 1), dtype=float16, numpy=
array([[0.],
       [0.]], dtype=float16)>)) is not compatible with B.Spec(b1=TensorSpec(shape=(2, 5), dtype=tf.float16, name=None), b2=TensorSpec(shape=(2, 6), dtype=tf.float32, name=None), a=A.Spec(a1=TensorSpec(shape=(2, 1), dtype=tf.float16, name=None))).


         [[{{node EagerPyFunc}}]] [Op:IteratorGetNext]

About this issue

  • Original URL
  • State: open
  • Created 2 years ago
  • Reactions: 5
  • Comments: 16 (7 by maintainers)

Most upvoted comments

I cannot reproduce this on TF 2.10.

Can you try to reproduce this on you system with TF 2.11.0-rc2?

+1
bhack on Nov 10, 2022
Read more comments on GitHub
← tensorflow: overparametrized convolution error in tf.nn.separable_conv2d
tensorflow: ParseTensor (tf.io.parse_tensor) is not vectorized - Vectorizing via tf.vectorized_map uses while_loop →