Improve generate docstring (for TF and FLAX)
# What does this PR do?
Just a continuation PR of https://github.com/huggingface/transformers/pull/18198 for TF and FLAX code

## Before submitting
- [x] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).

## Who can review?

@sgugger @gante 


_The documentation is not available anymore as the PR was closed or merged._<|||||>Looks like we just need a quick `make style` to be ready to merge :-)<|||||>Mmm no the formatting has done something very wrong here, we can't have that. There is likely some syntax error in the docstring.
Problem looks to be in the input_ids argument at line 422 of the TF generation files, the type should all be on one line.<|||||>> Mmm no the formatting has done something very wrong here, we can't have that. There is likely some syntax error in the docstring. Problem looks to be in the input_ids argument at line 422 of the TF generation files, the type should all be on one line.

`make extra_style_checks` does that automatically. What do you propose?<|||||>As I said, there is a syntax error in the docstring that makes the styling script behave erratically. The first step is to revert the changes, fix the syntax error then re-run it.<|||||>I changed this line
```
 input_ids (`tf.Tensor` of shape `(batch_size, sequence_length)`, `(batch_size, sequence_length,
            feature_dim)` or `(batch_size, num_channels, height, width)`, *optional*):
```
to this 
```
 input_ids (`tf.Tensor` of shape `(batch_size, sequence_length)`, `(batch_size, sequence_length, feature_dim)` or `(batch_size, num_channels, height, width)`, *optional*):
```
but then `make extra_style_checks` reverts the change <|||||>Ah yes, just tried locally and it's due to the empty line between `Parameters:` and `input_ids`. If you remove it, then your changes should not be overwritten.<|||||>> Ah yes, just tried locally and it's due to the empty line between `Parameters:` and `input_ids`. If you remove it, then your changes should not be overwritten.

Nice catch, but it still does look strange in the docs 🤔 

![image](https://user-images.githubusercontent.com/17574157/182453226-22ff1ea2-f620-43e5-958d-c19b1303878a.png)
<|||||>Uhmm, there is something wrong with the automatic styler -- e.g. the pytorch generate file should not be touched at all in this PR. As Sylvain wrote, the easiest solution is to start from a new branch 🤔 <|||||>> Uhmm, there is something wrong with the automatic styler -- e.g. the pytorch generate file should not be touched at all in this PR. As Sylvain wrote, the easiest solution is to start from a new branch 🤔

[The docs seem fine now](https://moon-ci-docs.huggingface.co/docs/transformers/pr_18432/en/main_classes/text_generation#transformers.generation_utils.GenerationMixin), right? It's just that this test, `doc-builder style src/transformers docs/source --max_len 119 --check_only --path_to_docs docs/source`, is not allowing to have more than 119 characters per line, but we need it here.<|||||>Nope, the docs are not fine for the PyTorch side with all the changes in this PR (and as @gante mentioned that file should not be touched at all). The doc-style is completely comfortable with lines that are more than 119 chars when it identifies they are parameter introduction lines, you just needed to remove the blank line between Parameters: and the first argument in `generate`.<|||||>Ah right, it has these strange hyphens...
![image](https://user-images.githubusercontent.com/17574157/182651706-dd36208c-d597-4d91-8233-02c385709b10.png)

I will close the PR then, let's disregard these changes.