XLA size-inference integration bug (tf.where, tf.TensorArray, slice, loop)
Hi @AndreasMadsen ,

I have replicated the reported behaviour and attached [gist](https://colab.research.google.com/gist/SuryanarayanaY/55cd78f6809a34a2ab9ea0d8dc2993c8/60742.ipynb) here. It needs to be dig into more for this behaviour. Meanwhile if you have already have some pointers and ready to submit a PR please let us know. Thanks @SuryanarayanaY, thanks for reproducing it and spending time on this. Unfortunately, I'm not familiar with the XLA code base and won't be making an attempt at solving this.

There are a few workarounds:

```diff
- x_batch = x_repeated[0:(batch_i + 1)*16, ...]
+ x_batch = x_repeated[0:tf.minimim(tf.shape(x_repeated)[0], (batch_i + 1)*16), ...]
```

and:

```diff
- predicts_all = predict_all_array.concat()
+ predicts_all = predict_all_array.concat()[:n_samples]
```

But this should not be necessary. 

PS: I suggest you add the `comp:xla` and `TF 2.12` labels. 
<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>