Add support of output_scores to flax models
# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Flax models does not support `output_scores` when generate() method is called,  despite the PyTorch models that fully supports this feature.
It is tried to follow naming and format of these parameters as same as PyTorch model codes (utils.py)

## Before submitting
- [x] This PR adds support of `output_scores` to flax models.
- [x] Flax Whisper model handles `output_scores` and `num_beams` parameters to consider during generate().

## Who can review?

@sanchit-gandhi 

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_22700). All of your documentation changes will be reflected on that endpoint.<|||||>@sanchit-gandhi @ArthurZucker 
Could you please review this PR?<|||||>I get a new error on CI for test codes of all models:
```AttributeError: module 'jax' has no attribute 'Array'```
However there was not such error two weeks ago (while my commit all tests passed).
Is there any updates on `Jax` that is not compatible with?
Any ideas @sanchit-gandhi ?

Update: The above mentioned error has been raised from `optax` source codes:
```
examples/flax/test_flax_examples.py:41: in <module>
    import run_clm_flax
examples/flax/language-modeling/run_clm_flax.py:40: in <module>
    import optax
../.pyenv/versions/3.8.12/lib/python3.8/site-packages/optax/__init__.py:18: in <module>
    from optax._src.alias import adabelief
../.pyenv/versions/3.8.12/lib/python3.8/site-packages/optax/_src/alias.py:23: in <module>
    from optax._src import clipping
../.pyenv/versions/3.8.12/lib/python3.8/site-packages/optax/_src/clipping.py:130: in <module>
    ) -> Tuple[List[chex.Array], jax.Array]:
E   AttributeError: module 'jax' has no attribute 'Array' 
```
Does it have relevance to this recent merge: https://github.com/huggingface/transformers/pull/22895 ?<|||||>> I get a new error on CI for test codes of all models: `AttributeError: module 'jax' has no attribute 'Array'` However there was not such error two weeks ago (while my commit all tests passed). Is there any updates on `Jax` that is not compatible with? Any ideas @sanchit-gandhi ?
> 
> Update: The above mentioned error has been raised from `optax` source codes:
> 
> ```
> examples/flax/test_flax_examples.py:41: in <module>
>     import run_clm_flax
> examples/flax/language-modeling/run_clm_flax.py:40: in <module>
>     import optax
> ../.pyenv/versions/3.8.12/lib/python3.8/site-packages/optax/__init__.py:18: in <module>
>     from optax._src.alias import adabelief
> ../.pyenv/versions/3.8.12/lib/python3.8/site-packages/optax/_src/alias.py:23: in <module>
>     from optax._src import clipping
> ../.pyenv/versions/3.8.12/lib/python3.8/site-packages/optax/_src/clipping.py:130: in <module>
>     ) -> Tuple[List[chex.Array], jax.Array]:
> E   AttributeError: module 'jax' has no attribute 'Array' 
> ```
> 
> Does it have relevance to this recent merge: #22895 ?

I've found that this issue was related to the optax version (which installed the 0.1.5). In the updated version of transformer repo, the version to be installed is forced to be 0.1.4<|||||>Good catch regarding the `jax.Array` issue! I need to un-pin JAX on Transformers since new Optax / Chex versions are running ahead https://github.com/huggingface/transformers/issues/19842 Will do this tomorrow 🤗<|||||>Thanks for the review @gante 🙌 See https://github.com/huggingface/transformers/pull/22700#discussion_r1182794360 for the next steps @hannan72 🚀<|||||>Also see related: #22700

This might get merged before this PR, in which case we can rebase to get the beam score fixes from main! Your changes will still be valuable for greedy search @hannan72 🤗<|||||>Hey @hannan72! This PR is looking in good shape - would you like to get it over the line with the last bits of integration?<|||||>> Hey @hannan72! This PR is looking in good shape - would you like to get it over the line with the last bits of integration?

Sorry for late response. I was busy on a product release.
Yes I really want to make it final and put it in the next release of transformers. What is remaining?
Please clarify the remaining steps to finalize the PR and close this issue.<|||||>Awesome! It's more or less as you left it - the major "TODO" is getting the correct vocab size in the first forward pass (see https://github.com/huggingface/transformers/pull/22700#discussion_r1178185868)<|||||>> Awesome! It's more or less as you left it - the major "TODO" is getting the correct vocab size in the first forward pass (see [#22700 (comment)](https://github.com/huggingface/transformers/pull/22700#discussion_r1178185868))

I had made a try on it and posted the result:


I tried to do this. But there was an error stopped me working on it.
I get the `vocab_size` from logits shape in the first step as follows:
```
next_tokens_scores = logits_processor(state.sequences, logits, state.cur_len)
next_token = jnp.argmax(next_tokens_scores, axis=-1)
scores = state.scores
if output_scores and state.scores is None:
    vocab_size = next_tokens_scores.shape[-1]
    scores = jnp.ones((batch_size, max_length, vocab_size)) * np.array(-1.0e7)
tokens_scores = scores.at[:, state.cur_len, :].set(next_tokens_scores) if output_scores else None
```

But in the line:
https://github.com/huggingface/transformers/blob/312b104ff65514736c0475814fec19e47425b0b5/src/transformers/generation/flax_utils.py#L641
 it checks that tensor shapes between runs should be exactly same, which causes the following error :
```
Exception has occurred: TypeError       (note: full exception trace is shown but execution is paused at: _run_module_as_main) body_fun output and input must have same type structure, got PyTreeDef(CustomNode(GreedyState[()], [*, *, *, *, *, {'decoder_attention_mask': *, 'decoder_position_ids': *, 'encoder_attention_mask': None, 'encoder_outputs': CustomNode(FlaxBaseModelOutput[()], [*, None, None]),...
```

So it seems the second suggestion is not going to work here. Because in Jax, every tensor shape should be pre-defined before deployment while we get the `vocab_size` during the deployment.<|||||>The idea here would be to run the first pass outside of the lax while loop (which we already do), then get the logits shape, then run the loop with the correct vocab size. Picking up on L730:
https://github.com/huggingface/transformers/blob/9ade58f0555430cec851e307c83c3a56c4a77d0b/src/transformers/generation/flax_utils.py#L730

This would look something like:

```python
        # The very first prompt often has sequence length > 1, so run outside of `lax.while_loop` to comply with TPU
        if input_ids.shape[1] > 1:
            state = sample_search_body_fn(state)

        # now get the vocab size
        vocab_size = state.logits.shape[-1]

        # do the other stuff that we need to do to init the state scores
        # ...

        # now run the main body
        if not trace:
            state = self._run_loop_in_debug(sample_search_cond_fn, sample_search_body_fn, state)
        else:
            state = lax.while_loop(sample_search_cond_fn, sample_search_body_fn, state)
```<|||||>> The idea here would be to run the first pass outside of the lax while loop (which we already do), then get the logits shape, then run the loop with the correct vocab size. Picking up on L730:
> 
> https://github.com/huggingface/transformers/blob/9ade58f0555430cec851e307c83c3a56c4a77d0b/src/transformers/generation/flax_utils.py#L730
> 
> This would look something like:
> 
> ```python
>         # The very first prompt often has sequence length > 1, so run outside of `lax.while_loop` to comply with TPU
>         if input_ids.shape[1] > 1:
>             state = sample_search_body_fn(state)
> 
>         # now get the vocab size
>         vocab_size = state.logits.shape[-1]
> 
>         # do the other stuff that we need to do to init the state scores
>         # ...
> 
>         # now run the main body
>         if not trace:
>             state = self._run_loop_in_debug(sample_search_cond_fn, sample_search_body_fn, state)
>         else:
>             state = lax.while_loop(sample_search_cond_fn, sample_search_body_fn, state)
> ```

I implemented your suggestion by applying following changes in `greedy_search_body_fn` and get the `vocab_size` from the first run as follows:

```
        def greedy_search_body_fn(state):
            """state update fn."""
            model_outputs = model(state.running_token, params=params, **state.model_kwargs)
            logits = model_outputs.logits[:, -1]

            # apply min_length, ...
            next_tokens_scores = logits_processor(state.sequences, logits, state.cur_len)

            next_token = jnp.argmax(next_tokens_scores, axis=-1)
            if output_scores:
                if state.scores is not None:
                    tokens_scores = state.scores.at[:, state.cur_len, :].set(next_tokens_scores)
                else:
                    scores = jnp.ones((batch_size, max_length, next_tokens_scores.shape[-1])) * np.array(-1.0e7)
                    tokens_scores = scores.at[:, state.cur_len, :].set(next_tokens_scores)
            else:
                tokens_scores = None
            next_token = next_token * ~state.is_sent_finished + pad_token_id * state.is_sent_finished
            next_is_sent_finished = state.is_sent_finished | (next_token == eos_token_id)
            next_token = next_token[:, None]

            next_sequences = lax.dynamic_update_slice(state.sequences, next_token, (0, state.cur_len))
            next_model_kwargs = self.update_inputs_for_generation(model_outputs, state.model_kwargs)
            return GreedyState(
                cur_len=state.cur_len + 1,
                sequences=next_sequences,
                scores=tokens_scores,
                running_token=next_token,
                is_sent_finished=next_is_sent_finished,
                model_kwargs=next_model_kwargs,
            )

        # The very first prompt often has sequence length > 1, so run outside of `lax.while_loop` to comply with TPU
        # Besides, when output_scores is true, to return scores vocab_size of the model is got from first run.
        if input_ids.shape[1] > 1 or output_scores:
            state = greedy_search_body_fn(state)

```<|||||>@sanchit-gandhi 
I think this PR is ready to merge. All tests are passed.
Could you please review it again and merge it?<|||||>@sanchit-gandhi 
I have checked out to my latest commit (b82ef360c5d819efc10298344d7d2fb4c33e1c47) and run a test as follows:

* Model_name: whisper_medium with flax inference
* GPU: A100-40GB
* Input audio: 5seconds
* transformers: git+https://github.com/huggingface/transformers.git@b82ef360c5d819efc10298344d7d2fb4c33e1c47
* Pytorch: 2.0.0
* jax: [cuda12_local] 0.4.11

A. Normal Inference (while `output_scores=False`):
The model has been deployed for 5 sequence runs. Inference time is ~0.2 seconds:

```
 model = FlaxWhisperForConditionalGeneration.from_pretrained(model_id, dtype=jnp.float16, from_pt=True)
 jit_generate = jax.jit(model.generate, static_argnames=["max_length", "language", "task"])
 runtime=[]
for i in range(5):
    start_time = time.time()
    input_features = jnp.array(input_features, dtype=jnp.float16)
    pred_ids = jit_generate(input_features, max_length=128, language='<|de|>', task ="transcribe")
    runtime.append(time.time() - start_time)
print("Inference time: ", runtime)
print("output scores: ", scores)
 ```
result:
Inference time: [57.01693844795227, 0.22632288932800293, 0.1981194019317627, 0.19892430305480957, 0.19736719131469727]
output scores: None


B. Inference with confidence scores (while `output_scores=True`):
The model has been deployed for 5 sequence runs. Inference time is also ~0.2 seconds:

```
 model = FlaxWhisperForConditionalGeneration.from_pretrained(model_id, dtype=jnp.float16, from_pt=True)
jit_generate = jax.jit(model.generate, static_argnames=["max_length", "language", "task", "output_hidden_states", "output_scores", "return_dict_in_generate"])
 runtime=[]
for i in range(5):
    start_time = time.time()
    input_features = jnp.array(input_features, dtype=jnp.float16)
    pred_ids = jit_generate(input_features, max_length=128, language='<|de|>', task ="transcribe",
                            output_scores=True, output_hidden_states=True, return_dict_in_generate=True)
    runtime.append(time.time() - start_time)
print("Inference time: ", runtime)
print("output scores: ", scores)
```
result:
Inference time: [82.8741066455841, 0.20504498481750488, 0.19746017456054688, 0.1972200870513916, 0.1973130702972412]
output scores: [[[-10000000. -10000000. -10000000. ... -10000000. -10000000. -10000000.]
  [      -inf       -inf       -inf ...       -inf       -inf       -inf]
  [      -inf       -inf       -inf ...       -inf       -inf       -inf]
  ...
  [-10000000. -10000000. -10000000. ... -10000000. -10000000. -10000000.]
  [-10000000. -10000000. -10000000. ... -10000000. -10000000. -10000000.]
  [-10000000. -10000000. -10000000. ... -10000000. -10000000. -10000000.]]]


It should be also noted that the result of model inference are exactly the same. The only change is that the first run takes more when `output_score=True` But next inferences are the approximately the same value.

Could you please review and merge this PR? 