[WIP] Flax pipeline support :pickup_truck: 
Author: Shubhamai
Pull request: yes

# What does this PR do?

Adds flax support in pipelines + corresponding flax changes of https://github.com/huggingface/transformers/pull/21516 by [ydshieh](https://github.com/ydshieh)

> Right before I was preparing to open pull request, with pretty much everything ready for review (including tests, ~ all tasks working), I was looking for issues this PR will fix, accidently, I found a PR [[WIP] Adding support for flax for pipelines.](https://github.com/huggingface/transformers/pull/14356/) (nearly 1-2 year old) by [Narsil](https://github.com/Narsil) previously unknown to me. I am adding this to alleviate any confusion, and after the commit `bug fixes & multiple tasks support`, I still continued with this PR since (I guess) a lot of `transformers` codebase has changed/updated, but I did used narsil's PR to add anything I missed in this PR, and to make it more polished. 

> So, I deeply thanks narsil and patrickvonplaten for the PR & review of the narsil's PR, as it helped this PR become more polished, and apologies as I should have checked if any similar PR is WIP.

*Not mentioning due to currently in WIP*

## Examples
```py
# Image Classifcation
vision_classifier = pipeline(task="image-classification", framework="flax", device="cuda:0") #cpu:0
print(vision_classifier(images="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"))

# Translation
en_fr_translator = pipeline("translation_en_to_fr", framework="flax")
print(en_fr_translator("How old are you?"))

# Image to Text
captioner = pipeline(model="ydshieh/vit-gpt2-coco-en", framework="flax")
print(captioner("https://huggingface.co/datasets/Narsil/image_dummy/raw/main/parrots.png"))
```

## Currently Supported tasks

Task	|PT|	TF|	Flax|	Reason for no Flax support|
|---|---|---|---|---|
audio_classification	|:heavy_check_mark:	|:x:|	:x:	|`AudioClassificationPipeline` is only available in PyTorch.|
automatic_speech_recognition|	:heavy_check_mark:|	:x:|	:x:	|The `AutomaticSpeechRecognitionPipeline` is only available in PyTorch.|
conversational	|:heavy_check_mark:	|:heavy_check_mark:|	:heavy_check_mark:|	
depth_estimation|	:heavy_check_mark:|	:x:|	:x:	|No Flax model available.|
document_question_answering|	:heavy_check_mark:|	:x:|	:x:|	No `FlaxAutoModelForDocumentQuestionAnswering` class available.|
feature_extraction	|:heavy_check_mark:|	:heavy_check_mark:	|:heavy_check_mark:|	
fill_mask|	:heavy_check_mark:|	:heavy_check_mark:|	:heavy_check_mark:	|
image_classification|	:heavy_check_mark:	|:heavy_check_mark:	|:heavy_check_mark:|	
image_segmentation|	:heavy_check_mark:|	:x:	|:x:	|No Flax model available.|
image_to_text|	:heavy_check_mark:	|:heavy_check_mark:	|:heavy_check_mark:	
object_detection|	:heavy_check_mark:|	:x:	|:x:|	No Flax model available.
question_answering|	:heavy_check_mark:|	:heavy_check_mark:|	:heavy_check_mark:	
summarization|	:heavy_check_mark:|	:heavy_check_mark:|	:heavy_check_mark:	
table_question_answering|	:heavy_check_mark:|	:heavy_check_mark:|	:x:|	No Flax model available.
text2text_generation|	:heavy_check_mark:|	:heavy_check_mark:|	:heavy_check_mark:	
text_classification	|:heavy_check_mark:|	:heavy_check_mark:|	:heavy_check_mark:	
text_generation	|:heavy_check_mark:	|:heavy_check_mark:	|:heavy_check_mark:	
token_classification|	:heavy_check_mark:|	:heavy_check_mark:|	:heavy_check_mark:	
translation	|:heavy_check_mark:	|:heavy_check_mark:|	:heavy_check_mark:	
video_classification|	:heavy_check_mark:	|:x:	|:x:|	No Flax model available.
visual_question_answering|	:heavy_check_mark:	|:x:|	:x:	|No Flax model available.
zero_shot_classification|	:heavy_check_mark:	|:heavy_check_mark:|	:heavy_check_mark:	
zero_shot_object_detection|	:heavy_check_mark:	|:x:|	:x:	|No Flax model available.

## Custom models link used in testing
~Could probably fix this by adding `from_pt=True` in `model_kwargs` in testing.~ Nope, during flax testing, PyTorch doesn't seems to be available. 
- [Shubhamai/distilbert-base-uncased-finetuned-sst-2-english](https://huggingface.co/Shubhamai/distilbert-base-uncased-finetuned-sst-2-english) < [distilbert-base-uncased-finetuned-sst-2-english](https://huggingface.co/distilbert-base-uncased-finetuned-sst-2-english)
- [Shubhamai/tiny-random-distilbert](https://huggingface.co/Shubhamai/tiny-random-distilbert) < [hf-internal-testing/tiny-random-distilbert](https://huggingface.co/hf-internal-testing/tiny-random-distilbert)
- [Shubhamai/tiny-random-vit](https://huggingface.co/Shubhamai/tiny-random-vit) < [hf-internal-testing/tiny-random-vit](https://huggingface.co/hf-internal-testing/tiny-random-vit)
- [Shubhamai/tiny-mbart](https://huggingface.co/Shubhamai/tiny-mbart) < [sshleifer/tiny-mbart](https://huggingface.co/)
- [Shubhamai/tiny-bert-for-token-classification](https://huggingface.co/Shubhamai/tiny-bert-for-token-classification) < [hf-internal-testing/tiny-bert-for-token-classification](https://huggingface.co/hf-internal-testing/tiny-bert-for-token-classification)
- [Shubhamai/tiny-random-clip-zero-shot-image-classification](https://huggingface.co/Shubhamai/tiny-random-clip-zero-shot-image-classification) < [hf-internal-testing/tiny-random-clip-zero-shot-image-classification](https://huggingface.co/hf-internal-testing/tiny-random-clip-zero-shot-image-classification)
- [Shubhamai/tiny-distilbert-base-cased-distilled-squad](https://huggingface.co/Shubhamai/tiny-distilbert-base-cased-distilled-squad) < [sshleifer/tiny-distilbert-base-cased-distilled-squad](https://huggingface.co/sshleifer/tiny-distilbert-base-cased-distilled-squad)


## Few questions for maintainers and users. 

- Should the framework name `pipeline(..., framework=)` be `jax` or `flax`, in this PR I used `flax` because we have been using them as a prefix in flax models, although we use `jax` as the alias in tokenizer functions, and in Huggingface Hub, so I am unable a make a final decision on this. 
    - If we use `framework="flax"`, a bit of a inconvenient code emerges, `self.tokenizer(inputs, return_tensors=self.framework if self.framework != "flax" else "jax")`, because tokenizer/image preprocessor doesn't identify `flax` framework, we have to change it to `jax`. Although it's a very small code, I get a bit uncomfortable presenting it, which could be a source of bugs. Solution would either rename `framework` to jax or add flax alias in tokenizer/image preprocessor, which function's exactly as jax.  
- JIT the model by default or another function argument?

Fixes 
- https://github.com/huggingface/transformers/issues/12627

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

- flax: sanchit-gandhi
- pipelines: Narsil

## TODO
 - [ ] Profiling & Benchmarking.
 - [ ] Updating docs.  

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_21851). All of your documentation changes will be reflected on that endpoint.

@Narsil , @sanchit-gandhi It looks good for an initial review. @sanchit-gandhi I am not also super confident if I have gotten the `jit` part right, I would appreciate if you can take a look and given any feedbacks on that. 

Thanks a lot for your time. 

Awesome to see this! Nice timing. automatic_speech_recognition would be cool to have too with `FlaxWhisper` that was merged recently!

Thanks for your PR @Shubhamai . At this stage, we do not have plans to have a pipeline in Flax similar to the ones in PyTorch and TensorFlow, which is why the original PR from @Narsil was not continued. There are several reasons for that:

1. The `pipeline` object is aimed at software engineers not necessarily familiar with machine learning, and Flax users are more researchers
2. `pipeline`s are meant to quickly try out a task, which does not work well in Flax/JAX where you have to compile and jit to get nice performance.

We can leave the PR open so that users try out your branch if they want something like a pipeline for Flax, but we don't want to commit to maintain this code, so we won't merge it in the main branch. What we might add in the future is a way to use large models on TPUs with Jax in a pipeline.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

> Thanks for your PR @Shubhamai . At this stage, we do not have plans to have a pipeline in Flax similar to the ones in PyTorch and TensorFlow, which is why the original PR from @Narsil was not continued. There are several reasons for that:
> 
> 1. The `pipeline` object is aimed at software engineers not necessarily familiar with machine learning, and Flax users are more researchers
> 2. `pipeline`s are meant to quickly try out a task, which does not work well in Flax/JAX where you have to compile and jit to get nice performance.
> 
> We can leave the PR open so that users try out your branch if they want something like a pipeline for Flax, but we don't want to commit to maintain this code, so we won't merge it in the main branch. What we might add in the future is a way to use large models on TPUs with Jax in a pipeline.

Hi,
These are very strange arguments against. I'm software engineer (or even someone like product owner/business user) and exactly because of that I want to use the same Pipeline API for different models! Flax model give the speed while inferring and why for their usage I need write a lot of  routine code (manually process inputs, manually generates and manually decode) instead of just one additional line for JIT compiling?

Particularly I'm interesting in `FlaxWhisperModel` and using it in `transformers.AutomaticSpeechRecognitionPipeline`


I would really recommend you the torch.compile way for Whisper as you can use it in a pipeline, and we have a lot more support for torch in general. Overall we don't have enough users requesting flax content + a growing torch ecosystem around the strong compile which pushed us a lot toward not continuing costly efforts in the flax domain!