Hannan Komari

@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...

> 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...

> 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....

> 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...

> 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...

@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 *...

Hello @sanchit-gandhi Sorry for late response I added test codes for output_scores feature of Flax that you asked me. But the PR is closed automatically by github-actions. My latest commit...

Thank you @sanchit-gandhi for re-opening this PR. I've add some tests based on what you asked me for. Could you please review my latest commit (https://github.com/huggingface/transformers/pull/22700/commits/0416becf86c65a3f32e72314715b79f5f84f52ce)?