Repository: tensorflow
Issue: XLA inference fails complaning about branch shape mismatches
Labels: stat:awaiting tensorflower,type:bug,comp:keras,comp:xla,TF 2.10

<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Bug

### Have you reproduced the bug with TF nightly?

Yes

### Source

binary

### Tensorflow Version

2.10.0

### Custom Code

Yes

### OS Platform and Distribution

Colab

### Mobile device

_No response_

### Python version

_No response_

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

A100 40 GB

### Current Behaviour?

```shell
I have the following function. 

First, I initialize a CLIP-based text encoder:


from keras_cv.models.stable_diffusion.text_encoder import TextEncoder

MAX_PROMPT_LENGTH = 77
text_encoder = TextEncoder(MAX_PROMPT_LENGTH)


Then, I am using the `text_encoder` like so in a function that I use to serialize the `text_encoder` as a `SavedModel`:

```python
from keras_cv.models.stable_diffusion.constants import _UNCONDITIONAL_TOKENS
import tensorflow as tf 

signature_dict = {
    "tokens": tf.TensorSpec(shape=[None, 77], dtype=tf.int32, name="tokens"),
}

def text_encoder_exporter(model: tf.keras.Model):
    BATCH_SIZE = 3
    MAX_PROMPT_LENGTH = 77
    POS_IDS = tf.convert_to_tensor([list(range(MAX_PROMPT_LENGTH))], dtype=tf.int32)
    UNCONDITIONAL_TOKENS = tf.convert_to_tensor([_UNCONDITIONAL_TOKENS], dtype=tf.int32)

    @tf.function(input_signature=[signature_dict])
    def serving_fn(inputs):
        # context
        encoded_text = model([inputs["tokens"], POS_IDS], training=False)
        encoded_text = tf.squeeze(encoded_text)

        if tf.rank(encoded_text) == 2:
            encoded_text = tf.repeat(
                tf.expand_dims(encoded_text, axis=0), BATCH_SIZE, axis=0
            )

        # unconditional context
        unconditional_context = model([UNCONDITIONAL_TOKENS, POS_IDS], training=False)

        unconditional_context = tf.repeat(unconditional_context, BATCH_SIZE, axis=0)
        return {"context": encoded_text, "unconditional_context": unconditional_context}

    return serving_fn
```

Serialization:

```python
tf.saved_model.save(
    text_encoder,
    "./text_encoder/1/",
    signatures={"serving_default": text_encoder_exporter(text_encoder)},
)
```

Now, while attempting to XLA-compile:

```python
from tensorflow.python.saved_model import tag_constants

batch_size = 3
saved_model_loaded = tf.saved_model.load(
    "./text_encoder/1/", tags=[tag_constants.SERVING]
)
text_encoder_predict_fn = saved_model_loaded.signatures["serving_default"]
# Raises error
xla_text_encoder_predict_fn = tf.function(text_encoder_predict_fn, jit_compile=True)
xla_text_encoder_predict_fn(
    tokens=tf.ones((batch_size, MAX_PROMPT_LENGTH), tf.int32)
).keys()
```
```


### Standalone code to reproduce the issue

```shell
https://colab.research.google.com/gist/sayakpaul/d7dafc252752a6c1ce10e85d8162b8ea/scratchpad.ipynb
```


### Relevant log output

```shell
/usr/local/lib/python3.8/dist-packages/tensorflow/python/util/traceback_utils.py in error_handler(*args, **kwargs)
    151     except Exception as e:
    152       filtered_tb = _process_traceback_frames(e.__traceback__)
--> 153       raise e.with_traceback(filtered_tb) from None
    154     finally:
    155       del filtered_tb

/usr/local/lib/python3.8/dist-packages/tensorflow/python/eager/execute.py in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     50   try:
     51     ctx.ensure_initialized()
---> 52     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
     53                                         inputs, attrs, num_outputs)
     54   except core._NotOkStatusException as e:

InvalidArgumentError: left_branch_shape.rank() != right_branch_shape.rank() (4 vs 3)
	 [[{{function_node __inference_serving_fn_55983}}{{node cond}}]] [Op:__inference_function_60000]
```
</details>