JAX/Flax models should be `jax.jit`ed by default? Or code examples should use jax.jit (~200x speedup)
Author: josephrocca
Pull request: no

I noticed that the Flax models were running really slow, and it took me a little while to realise that it was simply because they hadn't been `jit`ed. I'm new to JAX, so I could be missing something here, but wouldn't it make sense for `Flax<ModelName>.from_pretrained(...)` to return a pre-`jit`ed model?

If not, I wonder if it'd be a good idea to update the code examples so that newbies like me know that the model hasn't been `jit`ed?

## Environment info
- `transformers` version: 4.10.0
- Platform: Linux-5.11.0-7620-generic-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.6.9
- PyTorch version (GPU?): 1.10.0.dev20210622+cu111 (True)
- Tensorflow version (GPU?): 2.5.0 (True)
- Flax version (CPU?/GPU?/TPU?): 0.3.4 (gpu)
- Jax version: 0.2.17
- JaxLib version: 0.1.68
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No

### Who can help
@patil-suraj @patrickvonplaten

## Information

Model I am using: FlaxDistilBertForMaskedLM, FlaxCLIPModel, and other Flax models.

The problem arises when using the official example scripts.

## To reproduce

```python
!pip install --upgrade pip
!pip install transformers
!pip install --upgrade "jax[cuda111]" flax -f https://storage.googleapis.com/jax-releases/jax_releases.html

from transformers import DistilBertTokenizer, FlaxDistilBertForMaskedLM
import jax

tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
model = FlaxDistilBertForMaskedLM.from_pretrained('distilbert-base-uncased')

def test():
    inputs = tokenizer("The capital of France is [MASK].", return_tensors='jax')
    outputs = model(**inputs)
    logits = outputs.logits
    return logits
    
def test_jit():
    inputs = tokenizer("The capital of France is [MASK].", return_tensors='jax')
    outputs = jax.jit(model)(**inputs)
    logits = outputs.logits
    return logits

%timeit test()
# 623 ms ± 1.92 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit test_jit()
# 2.12 ms ± 329 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
```

```python
!pip install --upgrade pip
!pip install transformers Pillow requests
!pip install --upgrade "jax[cuda111]" flax -f https://storage.googleapis.com/jax-releases/jax_releases.html

import jax
from PIL import Image 
import requests
from transformers import CLIPProcessor, FlaxCLIPModel
model = FlaxCLIPModel.from_pretrained("openai/clip-vit-base-patch32")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

url = "http://images.cocodataset.org/val2017/000000039769.jpg"
image = Image.open(requests.get(url, stream=True).raw)

def test_jit():
    inputs = processor(text=["a photo of a cat", "a photo of a dog"], images=image, return_tensors="jax", padding=True)
    outputs = jax.jit(model)(**inputs)
    logits_per_image = outputs.logits_per_image
    probs = jax.nn.softmax(logits_per_image, axis=1)
    return probs

def test():
    inputs = processor(text=["a photo of a cat", "a photo of a dog"], images=image, return_tensors="jax", padding=True)
    outputs = model(**inputs)
    logits_per_image = outputs.logits_per_image
    probs = jax.nn.softmax(logits_per_image, axis=1)
    return probs

%timeit test()
# 2.5 s ± 58.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit test_jit()
# 15.4 ms ± 29.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
```

Comments:
Hi @josephrocca 

Thanks for the detailed issue. The Flax models are not jitted by default. This is by design, the user should be aware of such jax transformations like `jit`, `pmap`, `grad` etc, rather than the library abstracting it away. This gives you maximum control over what part of the code you want to `jit`. Also `jit` is for single device, for distributed training `pmap` is required, so the user should be aware of that, hence flax models are not transformed by default. And usually, one would `jit/pmap` as big a chunk of code as possible, for example the complete training step instead of just forward call. 

another reason is, the first call to jitted functions is usually slow which could, in turn, make the modeling tests really slow.

But you are right, the docs should be updated to make this clear and also provide examples of how to jit the model. Feel free to open a PR to update the docs if you are interested, happy to help with it :) 

@patil-suraj Ah, I see, thank you! Would it make sense for me to simply replace all instances of:
```python
outputs = model(**inputs)
```
with something like
```python
outputs = model(**inputs)  # or use jax.jit(model)(**inputs) for faster inference
```
for all the Flax model doc examples? Else could you suggest an appropriate way to edit the examples?

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/master/CONTRIBUTING.md) are likely to be ignored.