Tensor size-related conditional logic ignored in tf.function

[hosford42]

Click to expand!

Issue Type

Bug

Source

binary

Tensorflow Version

v2.9.0-18-gd8ce9f9c301

Custom Code

Yes

OS Platform and Distribution

Ubuntu 20.04.4 LTS

Mobile device

n/a

Python version

Python 3.8.10 (default, Jun 22 2022, 20:18:18) [GCC 9.4.0] on linux

Bazel version

n/a

GCC/Compiler version

n/a

CUDA/cuDNN version

No response

GPU model and memory

n/a

Current Behaviour?

I have a tf.function with a guard in place to make it return early under certain conditions, but the guard is being ignored, which causes an exception. I have tried rearranging the logic several different ways, but it always results in the same error.

Standalone code to reproduce the issue

import tensorflow as tf


@tf.function
def overwrite(arrays, indices, values) -> tf.Tensor:
    arrays = tf.convert_to_tensor(arrays)
    indices = tf.convert_to_tensor(indices)
    values = tf.convert_to_tensor(values)

    # No-op for empty arrays
    if tf.size(arrays) == 0:
        return arrays

    indices = indices % tf.shape(arrays)[1]
    arrays = tf.tensor_scatter_nd_update(
        arrays,
        tf.concat([tf.range(tf.shape(arrays)[0])[:, tf.newaxis], indices[:, tf.newaxis]], axis=-1),
        values
    )

    return arrays


print()
print("Works fine:")
print(overwrite(tf.zeros((4, 2)), tf.range(4), tf.ones((4,))))
print()
print("Causes an exception when it ought to exit early:")
print(overwrite(tf.zeros((4, 0)), tf.range(4), tf.ones((4,))))
print()

Relevant log output

2022-08-28 20:49:57.080725: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:975] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
...
2022-08-28 20:49:57.239798: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
...
2022-08-28 20:49:59.004862: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1532] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 3359 MB memory:  -> device: 0, name: NVIDIA GeForce GTX 1050 Ti, pci bus id: 0000:01:00.0, compute capability: 6.1

Works fine:
tf.Tensor(
[[1. 0.]
 [0. 1.]
 [1. 0.]
 [0. 1.]], shape=(4, 2), dtype=float32)

Causes an exception when it ought to exit early:
Traceback (most recent call last):
  File "/usr/lib/python3/dist-packages/IPython/core/interactiveshell.py", line 3331, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-7-0c355e479cbf>", line 1, in <module>
    print(overwrite(tf.zeros((4, 0)), tf.range(4), tf.ones((4,))))
  File "/home/hosford42/.local/lib/python3.8/site-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/tmp/__autograph_generated_file7d47e8wb.py", line 40, in tf__overwrite
    ag__.if_stmt((ag__.converted_call(ag__.ld(tf).size, (ag__.ld(arrays),), None, fscope) == 0), if_body, else_body, get_state, set_state, ('do_return', 'retval_', 'arrays', 'indices'), 2)
  File "/tmp/__autograph_generated_file7d47e8wb.py", line 33, in else_body
    arrays = ag__.converted_call(ag__.ld(tf).tensor_scatter_nd_update, (ag__.ld(arrays), ag__.converted_call(ag__.ld(tf).concat, ([ag__.converted_call(ag__.ld(tf).range, (ag__.converted_call(ag__.ld(tf).shape, (ag__.ld(arrays),), None, fscope)[0],), None, fscope)[:, ag__.ld(tf).newaxis], ag__.ld(indices)[:, ag__.ld(tf).newaxis]],), dict(axis=(- 1)), fscope), ag__.ld(values)), None, fscope)
ValueError: in user code:
    File "<ipython-input-3-dcfd69f4f71e>", line 12, in overwrite  *
        arrays = tf.tensor_scatter_nd_update(
    ValueError: Indices and updates specified for empty input for '{{node cond/TensorScatterUpdate}} = TensorScatterUpdate[T=DT_FLOAT, Tindices=DT_INT32](cond/TensorScatterUpdate/arrays, cond/concat, cond/TensorScatterUpdate/values)' with input shapes: [4,0], [4,2], [4].

[hosford42]

Here's an even simpler way to reproduce the issue, which doesn't involve tensor_scatter_nd_update at all:

import tensorflow as tf

@tf.function
def f(x):
    x = tf.convert_to_tensor(x)
    if tf.size(x) > 0:
        tf.assert_greater(tf.size(x), 0)
    return x

f(tf.zeros((1,)))  # Works fine
f(tf.zeros((0,)))  # Bombs out

Here's the traceback:

Traceback (most recent call last):
  File "/usr/lib/python3/dist-packages/IPython/core/interactiveshell.py", line 3331, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-8-6c94edddf473>", line 1, in <module>
    f(tf.zeros((0,)))
  File "/home/hosford42/.local/lib/python3.8/site-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/home/hosford42/.local/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py", line 1127, in autograph_handler
    raise e.ag_error_metadata.to_exception(e)
tensorflow.python.framework.errors_impl.InvalidArgumentError: in user code:
    File "<ipython-input-6-df56e54e761f>", line 5, in f  *
        tf.assert_greater(tf.size(x), 0)
    InvalidArgumentError: Condition x > y did not hold element-wise:
    x (cond/Size:0) = 
    0
    y (cond/assert_greater/y:0) = 
    0

[hosford42]

I even tried building a tf.cond manually, to no avail:

import tensorflow as tf


@tf.function
def f(x):
    x = tf.convert_to_tensor(x)
    return tf.cond(tf.size(x) == 0, lambda: 0, lambda: (tf.assert_greater(tf.size(x), 0), tf.size(x))[-1])


f(tf.zeros((1,)))
f(tf.zeros((0,)))

Traceback:

Traceback (most recent call last):
  File "/usr/lib/python3/dist-packages/IPython/core/interactiveshell.py", line 3331, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-3-24fccfd78995>", line 11, in <module>
    f(tf.zeros((0,)))
  File "/home/hosford42/.local/lib/python3.8/site-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/home/hosford42/.local/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py", line 1127, in autograph_handler
    raise e.ag_error_metadata.to_exception(e)
tensorflow.python.framework.errors_impl.InvalidArgumentError: in user code:
    File "<ipython-input-3-24fccfd78995>", line 7, in f  *
        return tf.cond(tf.size(x) == 0, lambda: 0, lambda: (tf.assert_greater(tf.size(x), 0), tf.size(x))[-1])
    InvalidArgumentError: Condition x > y did not hold element-wise:
    x (cond/Size:0) = 
    0
    y (cond/assert_greater/y:0) = 
    0

[hosford42]

At this point, I see no other options for working around this besides modifying the tensor to be non-empty, and then cutting it back down to size again afterwards.

[tilakrayal]

@gadagashwini,
I was able to reproduce the issue on tensorflow v2.8, v2.9 and nightly. Kindly find the gist of it here.