XLA compiled tf.where with known output shape error (set_shape, tf.concat / tf.stack)
can i work on this ? @mirjanic Could you please have a look at this [gist](https://colab.research.google.com/gist/sushreebarsa/e3b2bf759b7c10f7ca7ba32fb0bb6abd/untitled0.ipynb#scrollTo=c7xn4WUH5RH4) and confirm this issue?
The error log is as follows;
```
Shape before (unknown): (10, 2)
Shape after (known): (10, 2)
Shape before (unknown): (None, 2)
Shape after (known): (10, 2)
Shape before (unknown): (None, 2)
Shape after (known): (10, 2)
---------------------------------------------------------------------------
InvalidArgumentError                      
```

Thank you! Yes, this is the issue I am having. The error message "InvalidArgumentError: Cannot concatenate arrays that differ in dimensions other than the one being concatenated" occurs when trying to concatenate two tensors with different shapes


In the given code, the concatenation fails because the shapes of x and y are different. Specifically, x has shape (10, 2) after its shape is set, while y has shape (10, 1) 

To fix this error, the shapes of x and y need to be made compatible for concatenation. One possible solution is to change the shape of y to (10, 2) by replacing tf.expand_dims(tf.range(x.shape * 2, dtype=tf.int64), axis=-1) with tf.range(x.shape * 2, dtype=tf.int64)[:, tf.newaxis] 

This will create a tensor with shape (10, 2) that can be concatenated with x.

Here is the update code . 

https://colab.research.google.com/drive/1rS1T-T7yKH6VyF1bp2st5PLT0FQ-L56g?usp=sharing

Is it okay now ? Unfortunately, I do not follow your suggestion. You mention that shapes `(10,2)` and `(10,1)` are incompatible for concatenation, but they are; with the result having shape `(10,3)`. You propose reshaping `y` to `(10,2)`, but you do not do that: in your code `y` still has shape `(10,1)` and the concatenation works just as well (as it should).

The real issue here is that the snippet in OP fails _when compiled with XLA_ because compiler does not understand the true shape of `x` even after `set_shape` is called and correct length is updated from `None` to `10`.

According to the documentation, `set_shape` is supposed to "provide additional shape information that cannot be inferred from the graph alone", but this information does not seem to be used correctly by XLA.

I hope this clarifies the issue! Try with `tf.RaggedTensor` instead of using `tf.concat` for concatenation. `tf.RaggedTensor` can handle variable-length dimensions and accommodate the shape mismatch that occurs when concatenating tensors with different dimensions. Convert the `tf.where` output to `tf.RaggedTensor` and then use the appropriate method provided by `tf.RaggedTensor` to concatenate the tensors. I thought Ragged Tensors were [unsupported in XLA](https://github.com/tensorflow/tensorflow/issues/56595)? Also, these tensors have matching shapes and the code above runs both with eager and with graph execution. It only fails _in XLA_ because it is not using the static shape that I am trying to provide to it. Here is a distilled version of the problem ([Colab URL](https://colab.research.google.com/drive/1nLIPcrT3Zz5MZ78D7Blhcj-PTeqIlZXn?usp=sharing)):

```python
import tensorflow as tf
import numpy as np

# Arbitrary parameters
n=10
k=4

def fun(x):
    # print(x.shape)
    # x.shape = (n,)

    y = tf.one_hot(x, k)
    # print(y.shape)
    # y.shape = (n, k)

    z = tf.where(y == 1)
    # print(z.shape)
    # z.shape = (None, 2)

    # z.shape cannot be determined statically due to
    # tf.where call, but we know that it must be
    # equal to (n, 2), so we try to set it manually:

    # Variant 1:
    z.set_shape(shape=(n,2))
    # print(z.shape)
    # z.shape = (n, 2)
    # We have used set_shape to supply additional info
    # about the shape that tf could not deduce, and we see
    # that the shape is updated.

    # Variant 2:
    # z = tf.ensure_shape(z, shape=(n, 2))  # <-- FAILS!
    # Interestingly, using ensure_shape fails immediately,
    # with the message that z.shape is (n*k, 2), not (n, 2)
    # This is perhaps the core problem because this violates
    # tf.where docs which say that "The result's shape is
    # [tf.math.count_nonzero(condition), tf.rank(condition)]".

    tf.concat([x[:, tf.newaxis], z], axis=1)  # Fails only in XLA

x = np.random.randint(k, size=n)
x = tf.constant(x, dtype=tf.int64)

fun(x)  # Succeeds
tf.function(fun)(x)  # Succeeds
tf.function(fun, jit_compile=True)(x)  # Fails!
```

Error message produced by `tf.ensure_shape` suggests that output shape of `tf.where` is wrong when using XLA. However, if we define `x` inside `fun` (as opposed to passing an argument), then `fun` compiles and runs successfully, which further muddies the water. Closing because I found a combination of `autoclustering` and `tf.where` replacement that circumvents the issue. Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/61755">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/61755">No</a>

### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

No

### Source

binary

### TensorFlow version

2.13

### Custom code

No

### OS platform and distribution

Google Colab

### Mobile device

_No response_

### Python version

3.10.12

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

Output of `tf.where`, when used inside `tf.function` with `jit_compile=True`, can sometimes be used correctly (as with sum), and sometimes raises shape mismatch error (as with concatenation). This error is present even if output shape is set manually with `set_shape`. 

The code below runs without `jit_compile` or with sum instead of `tf.concat`, and only fails if concatenating inside a compiled function. 

Note: `autoclustering` solves the issue on the toy example, but _not_ on the codebase I am working on.

### Standalone code to reproduce the issue

Colab: https://colab.research.google.com/drive/1FuboVMSao8eCZLcZ2F7Fdsa1UvlnHQmG?usp=sharing

```python
import tensorflow as tf

def fun(x, y):
    x = tf.where(x == 1)
    print(f'Shape before (unknown): {x.shape}') 
    x.set_shape(shape=[y.shape[0], 2])
    print(f'Shape after (known): {x.shape}') 
    return tf.concat([x,y], axis=1)  # Concatentation fails
    # return x + y  # Sum would succeed

x = tf.constant([[0,0,1,1,0],
                  [0,1,0,1,0],
                  [1,0,0,0,1],
                  [1,0,1,0,0],
                  [0,1,1,0,0],], dtype=tf.int32)
y = tf.expand_dims(tf.range(x.shape[0] * 2, dtype=tf.int64), axis=-1)

fun(x, y)

tf.function(fun)(x,y)

tf.function(fun, jit_compile=True)(x,y)  # Fails as described above
```


### Relevant log output

```shell
Shape before (unknown): (None, 2)
Shape after (known): (10, 2)
InvalidArgumentError: Cannot concatenate arrays that differ in dimensions other than the one being concatenated. Dimension 0 in both shapes must be equal: s64[<=25,2] vs s64[10,1].
```
