Repository: tensorflow
Issue: tf.concat (and tf.transpose) inside a for loop with tf.range in the context of a GradientTape while using XLA dosn't work
Labels: stat:awaiting tensorflower,type:bug,comp:xla,TF 2.15

### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

No

### Source

binary

### TensorFlow version

2.15.0.post1, 2.15.0

### Custom code

Yes

### OS platform and distribution

WSL Ubuntu 22.04

### Mobile device

_No response_

### Python version

3.10.12

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

12.2

### GPU model and memory

RTX 3080 Ti

### Current behavior?

### Current behavior
Even tho the shapes of the elements which are concatenated are well defined to compile time, when running the code provided, the execution fails with the error log provided. This is also the case, if you enforce the shapes with using tf.reshape(). 
If you remove either 
- the GradientTape context
- or the for loop with tf.range
- or the XLA compilation 

the code will work as expected.

Unrolling the for loop with range() is not desired since the number of iterations will be >50000 in the final project. Also converting the for loop in a tf.while loop with maximum_iterations specified will result in the same error. Specifying the batchsize to a constant value (also in the input_signature) won't resolve the issue either. Also changing the tf.device between GPU / CPU won't resolve the issue.
The same error arises if you try to use tf.transpose()


### Expected behavior
The arrays with well defined shapes at compile time will be concatenated / transposed when using a for loop with tf.range in the context of a GradientTape while using XLA (jit_compile=True)

### Standalone code to reproduce the issue

```shell
class Model():
    def __init__(self, batchsize):
        self.batchsize = batchsize
        self.g = tf.Variable(initial_value=tf.ones([self.batchsize,1,1]), trainable=True, name="g")
        self.m = tf.Variable(initial_value=tf.ones([self.batchsize,1,1]), trainable=True, name="m")
        self.d = tf.Variable(initial_value=tf.ones([self.batchsize,1,1]), trainable=True, name="d")
        self.k = tf.Variable(initial_value=tf.ones([self.batchsize,1,1]), trainable=True, name="k")
        self.y0 = tf.ones([self.batchsize, 2, 1])


    @tf.function(input_signature=[tf.TensorSpec(shape=(None, 2, 1), dtype=tf.float32)],
                 jit_compile=False)
    def compute(self, yi):
        A1 = tf.concat([tf.zeros_like(self.m), tf.ones_like(self.m)], 2)
        A2 = tf.concat([-self.k/self.m, -self.d/self.m], 2)
        A = tf.concat([A1, A2], 1)
        #A = tf.transpose(A, perm=[0, 2, 1])
        B = tf.concat([tf.zeros_like(self.g), self.g], 1)

        dy = tf.linalg.matmul(A, yi) + B
        return dy
    

class Estimator():
    def __init__(self, model):
        self.model = model
        self.dt = tf.constant(0.001)

    @tf.function(jit_compile=True, reduce_retracing=False)
    def estimate(self):
        with tf.GradientTape() as tape:
            yi = self.model.y0
            for i in tf.range(10):
                dyi = self.model.compute(yi)
                yi = yi + dyi*self.dt
        grads = tape.gradient(yi, [self.model.g, self.model.m, self.model.d, self.model.k])
        return grads



device = "CPU:0"
#device = "GPU:0"

with tf.device(device):    
    batchsize = 5
    model = Model(batchsize)
    estimator = Estimator(model)
    grads = estimator.estimate()
```


### Relevant log output

```shell
OP_REQUIRES failed at concat_op.cc:168 : INVALID_ARGUMENT: Input 0 to node `gradient_tape/while/gradients/while/StatefulPartitionedCall_grad/PartitionedCall/gradients/concat_3_grad/ConcatOffset` with op ConcatOffset must be a compile-time constant.

XLA compilation requires that operator arguments that represent shapes or dimensions be evaluated to concrete values at compile time. This error means that a shape or dimension argument could not be evaluated at compile time, usually because the value of the argument depends on a parameter to the computation, on a variable, or on a stateful operation such as a random number generator.
```