TensorArray issues in XLA

[yongwww]

Please make sure that this is a bug. As per our
GitHub Policy,
we only address code/doc bugs, performance issues, feature requests and
build/installation issues on GitHub. tag:bug_template

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): No
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Ubuntu 16.04
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: NA
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): 2.3 (exists in 1.5 as well)
• Python version: 3.7.6
• Bazel version (if compiling from source): NA
• GCC/Compiler version (if compiling from source): NA
• CUDA/cuDNN version: NA
• GPU model and memory: NA

You can collect some of this information using our environment capture
script
You can also obtain the TensorFlow version with:

1. TF 1.0: python -c "import tensorflow as tf; print(tf.GIT_VERSION, tf.VERSION)"
2. TF 2.0: python -c "import tensorflow as tf; print(tf.version.GIT_VERSION, tf.version.VERSION)"

Describe the current behavior
I know it was documented in xla official doc that xla doesn't support TensorArray with dynamic_size=True, however, I still ran into issues even w/ dynamic_size=False(it is default).

Describe the expected behavior
expect the tensorarray related stuffs are supported.

Standalone code to reproduce the issue
Provide a reproducible test case that is the bare minimum necessary to generate
the problem. If possible, please share a link to Colab/Jupyter/any notebook.

import numpy as np
import tensorflow as tf
tf.compat.v1.enable_eager_execution()

@tf.function(experimental_compile=True)
def test_tensor_array_scatter_gather():
    dtype = "float32"
    t = tf.constant(np.array([[1.0], [2.0], [3.0]]).astype(dtype))
    scatter_indices = tf.constant([2, 1, 0])
    gather_indices = tf.constant([1, 2])
    ta1 = tf.TensorArray(dtype=dtype, size=3, infer_shape=True)
    ta2 = ta1.scatter(scatter_indices, t)
    t1 = ta2.gather(gather_indices)
    return t1

test_tensor_array_scatter_gather()

Other info / logs Include any logs or source code that would be helpful to
diagnose the problem. If including tracebacks, please include the full
traceback. Large logs and files should be attached.

[ravikyram]

@yongwww

I have tried in colab with TF versions 2.3, nightly versions(2.4.0-dev20200831) and was able to reproduce the issue.If I change
@tf.function(experimental_compile=True) to @tf.function i am not seeing any issue.Looks like it does not utilize XLA.Please, find the gist here..Thanks!

[yongwww]

@ravikyram yeah, I think @tf.function doesn't use XLA either.

[yongwww]

I have several other tensorarray test cases, will run into similar issues with all of them, I can upload some of them if needed.

[sachinprasadhs]

Was able to reproduce your issue in TF 2.11, please find the gist here. Thanks!