Repository: tensorflow
Issue: XLA size-inference integration bug (tf.where, tf.TensorArray, slice, loop)
Labels: stat:awaiting tensorflower,type:bug,comp:xla,TF 2.12

<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Bug

### Have you reproduced the bug with TF nightly?

Yes

### Source

binary

### Tensorflow Version

2.14.0.dev20230502

### Custom Code

No

### OS Platform and Distribution

Google Colab

### Mobile device

_No response_

### Python version

3.10.11

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current Behaviour?

The attached code produces difference results for `jit_compile=True` and `jit_compile=False`, the XLA version produces the wrong results. The bug reproduces on both TF 2.12 and nightly (2.14.0.dev20230502) on Google Colab.

It appears to be an issue with `tf.size(tf.where(x == 1))`. While the visible front-end output is correct, I suspect there is an integration bug on the compiler back-end. When combined with a slice that is too big x[0:16] (`x.shape = (3, 5)`), plus many more code details, the output has same size as the input (`5` in the example). As if `tf.size(tf.where(x == 1)) == tf.size(x) == 5`, when actually `tf.size(tf.where(x == 1)) == 2`. I suspect this because changing the input shape to `(4, )` changes the output shape to `(4, )`.

This bug was very subtle and difficult to find, especially because `tf.print()` is not supported by XLA. I hope you appreciate it.

### Standalone code to reproduce the issue

Google Colab link: https://colab.research.google.com/drive/1cccEssWkCVLO515uf3RoOuEMDp5V-2ZZ?usp=sharing

```python
import tensorflow as tf

def fn(x, y):
    token_idx_to_mask = tf.where(x == 1)
    n_samples = tf.size(token_idx_to_mask)
    x_repeated = tf.repeat(tf.expand_dims(x, 0), n_samples, axis=0)

    predict_all_array = tf.TensorArray(x_repeated.dtype, size=1, infer_shape=False, element_shape=(None, ))
    for batch_i in tf.range(1):
        x_batch = x_repeated[0:(batch_i + 1)*16, ...]
        y_batch = x_batch[:, y]
        predict_all_array = predict_all_array.write(batch_i, y_batch)
    predicts_all = predict_all_array.concat()

    return n_samples, predicts_all

fn_jit = tf.function(reduce_retracing=True, jit_compile=True)(fn)
fn_std = tf.function(reduce_retracing=True)(fn)

print('XLA': fn_jit(tf.constant([0, 1, 1, 0, 0], dtype=tf.dtypes.int32), tf.constant(0)))
print('STD': fn_std(tf.constant([0, 1, 1, 0, 0], dtype=tf.dtypes.int32), tf.constant(0)))
```


### Relevant log output

```shell
XLA: (<tf.Tensor: shape=(), dtype=int32, numpy=2>, <tf.Tensor: shape=(5,), dtype=int32, numpy=array([0, 0, 0, 0, 0], dtype=int32)>)
STD: (<tf.Tensor: shape=(), dtype=int32, numpy=2>, <tf.Tensor: shape=(2,), dtype=int32, numpy=array([0, 0], dtype=int32)>)
```

Expected results would be:

```shell
XLA: (<tf.Tensor: shape=(), dtype=int32, numpy=2>, <tf.Tensor: shape=(2,), dtype=int32, numpy=array([0, 0], dtype=int32)>)
STD: (<tf.Tensor: shape=(), dtype=int32, numpy=2>, <tf.Tensor: shape=(2,), dtype=int32, numpy=array([0, 0], dtype=int32)>)
```
</details>