Add Flax image captioning example
# What does this PR do?

Add `run_image_captioning_flax.py` (modified from `run_summarization_flax.py`).

## Who can review

Examples:  @patil-suraj + cc @patrickvonplaten @NielsRogge @sgugger for info

@patil-suraj , @sgugger 

Thank you for the comments. I will make the example much simpler by only supporting loading pre-trained vision encoder and text models.

About casting jax array to np arrary, there is a significant slow down (at least, in image examples) when using jax array as indices for accessing `datasets.Dataset`. I will re-produce the timing comparison.


<|||||>Hi, @patil-suraj @sgugger 

I simplified the config/model initialization parts (only support loading pretrained encoder & decoder).

--------------------

For @patil-suraj 

About using `numpy` array instead of `jnp.array` when it comes to `datasets`,

For this line

https://github.com/huggingface/transformers/blob/650fb4aa2adfe746f8c3e1aec1d601c4a0e9f40c/examples/flax/image-captioning/run_image_captioning_flax.py#L854

takes 30 seconds (for selecting `16384` elements) using `jax.numpy`,  while using `numpy` only takes `0.005` second.

For this line (take 256 elements - with image data)

https://github.com/huggingface/transformers/blob/650fb4aa2adfe746f8c3e1aec1d601c4a0e9f40c/examples/flax/image-captioning/run_image_captioning_flax.py#L348

jax.numpy: 0.45 second / numpy: 0.10 - 0.15 second

A singe training step (global batch size: 256 images) takes < 0.5 seconds on TPU. 

Due to this significant differences in processing speed, I think it is worth keeping using `numpy` when dealing with `datasets`.
Let me know if you have different opinions about this :-)
<|||||>Hi

- support `from_pretrained` (the model creation from encoder/decoder is done in another script)
- README updated
- rename to `coco_dataset_script` to avoid confusion
- other nits applied

Thanks for the reviews :-)