Add `output_hidden_state` and `output_scores` to Flax generate

I need whisper's output_scores and output_hidden_states as the result of generate() method. On Pytorch model, I can easily get the output_scores and output_hidden_states by setting these parameters in generate() method as follows: ``` whisper_output = model.generate(inputs=input_features, max_new_tokens=180, output_scores=True, output_hidden_states=True, return_dict_in_generate=True) ``` and the resulted `whisper_output` returns 'scores' and 'output_hidden_states' as it keys alongside 'sequences' Now I want to do so for Flax whisper model. but setting these parameters as the static_argnames of model doesn't have effect to get output_scores. Is there any solution for getting output_scores or logits from Flax whisper model?