(TF) model.generate with beam search to an exportable tf.function
### Feature request

It would be good if we could export the `tf.function` `model.generate` with beam search, so it is servable with TF Serving. 

### Motivation

When we try to export a `model.generate` with `do_sample=True` to a `tf.function` and save it, we get an error due to Transformers and Tensorflow Internals. It would be good if, such as what is done in #18372, we could do for beam search:

```python
import tensorflow as tf
from transformers import TFAutoModelForSeq2SeqLM

class MyOwnModel(tf.Module):
    def __init__(self, model_path="t5-small"):
        super(MyOwnModel, self).__init__()
        self.model = TFAutoModelForSeq2SeqLM.from_pretrained(model_path)

    @tf.function(input_signature=(tf.TensorSpec((None, 32), tf.int32, name="input_ids"),
                                  tf.TensorSpec((None, 32), tf.int32, name="attention_mask")), jit_compile=True)
    def serving(self, input_ids, attention_mask):
        outputs = self.model.generate(input_ids=input_ids, attention_mask=attention_mask, max_new_tokens=32,
                                   return_dict_in_generate=True, do_sample=True)
        return {"sequences": outputs["sequences"]}

model = MyOwnModel()
export_dir = "./t5-example"

tf.saved_model.save(
    model,
    export_dir,
    signatures={
        "serving_default":
            model.serving
    })
```

That way, we would be able to serve TF text to text models with beam search in TF Serving, which would be a blast. Otherwise we would have to implement our own servers which is not ideal.

### Your contribution

I can try submitting a PR, but would be glad if I could get some guidance because this one seems harsh. It seems like we would have to remove lots of `shape` related operations. 
cc @gante <|||||>Hi @piEsposito 👋 Thanks for raising this issue -- there were indeed two related problems, both fixed in #19773 :)<|||||>@gante you always bring good news when I'm here to bother you folks from Hugging Face. Thanks!