Flax model much slower than Pytorch
Hi,

I am running the same QuestionAnswering model using Pytorch and Flax, but I am seeing quite a bad performance for Flax.

| Model         	| Time   	|
|---------------	|--------	|
| Pytorch (GPU) 	| 30 ms  	|
| Flax (GPU)    	| 4 sec  	|
| Flax (TPU)    	| 15 sec 	|

I have warmed the jit function, and I think I have put the model and data on the device in both cases. What am I doing wrong here?

Notebook: https://colab.research.google.com/drive/18ouY4S9rYtlEC1dht_w8WU4fFKnP9yhP#scrollTo=FLfV3ymlwRJC


## Environment info
- `transformers` version: 4.16.2
- Platform: Linux-5.4.144+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.7.12
- PyTorch version (GPU?): 1.10.0+cu111 (True)
- Tensorflow version (GPU?): 2.7.0 (True)
- Flax version (CPU?/GPU?/TPU?): 0.4.0 (gpu)
- Jax version: 0.2.25
- JaxLib version: 0.1.71
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No

### Who can help
@patrickvonplaten 

## Information

Model I am using (Bert, XLNet ...): BERT (FlaxBertForQuestionAnswering, BertForQuestionAnswering)

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [ X] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [ X] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

1. Run [this notebook](https://colab.research.google.com/drive/18ouY4S9rYtlEC1dht_w8WU4fFKnP9yhP#scrollTo=FLfV3ymlwRJC) on Google Colab (GPU runtime)

## Expected behavior
About the same performance for Pytorch as for Flax.

cc @patil-suraj think we need to improve our docs here. 

@yorickvanzweeden the very first time you run `run_model_jit(jax_inputs_batched)`, the model is compiled which takes time. However the second time you run it, it should be pretty fast. Could you maybe run the code say 20 times for both PyTorch and Flax and then check if the Flax is still as slow?<|||||>@patrickvonplaten Thank you for your swift response!

I noticed caching on the inputs, so I have two inputs. When I do something like this, I still get about 4 seconds versus 40 ms.

```python
# PyTorch
%timeit -r 20 -n 1 model(**inputs_gpu) # <-- 20x repeat
> 1 loop, best of 20: 20.5 ms per loop
%timeit -r 1 -n 1 model(**inputs_gpu2)
> 1 loop, best of 1: 39.1 ms per loop

# Flax
%timeit -r 20 -n 1 run_model_jit(jax_inputs_batched) # <-- 20x repeat
> 1 loop, best of 20: 44.7 µs per loop
%timeit -r 1 -n 1 run_model_jit(jax_inputs_batched2)
> 1 loop, best of 1: 3.83 s per loop 
```

Also, please note that I am using `FlaxBertForQuestionAnsweringModule`, and not `FlaxBertForQuestionAnswering` directly. I would like to use the module directly.

But also when using `FlaxBertForQuestionAnswering.from_pretrained`, I get the second input in seconds (5.75 sec).

I updated [the notebook](https://colab.research.google.com/drive/18ouY4S9rYtlEC1dht_w8WU4fFKnP9yhP#scrollTo=d5Ic_18Affx0) for the 20x repeat, and using the model rather than the module.

<|||||>Hey @yorickvanzweeden, 

Cool, so it seems that once the flax inputs are cached, the forward pass is very fast no? 

Now the problem is that for the second input `jax_inputs_batched2` the code is recompiled again. The most important thing when dealing with `jax.jit(...)` is to remember the following if you `jax.jit(...)` to create a "jitted" function, *e.g.* `run_model_jit(...)` and then run the function, *e.g.* `run_model_jit(jax_inputs_batched2)`, the following happens under the hood:

1. JAX checks if it can find a compiled binary that corresponds to the function you jitted `run_model_jit(...)` **and** the input shape you forwarded.
2. If there is no match, then JAX compiles the function + input shape and we repeat step 1. If there is a match then the compiled binary is executed.

In your case for `run_model_jit(jax_inputs_batched2)` JAX does find a compiled binary, but the compiled binary corresponds to `run_model_jit` + `jax_inputs_batched.shape` and while you again use `run_model_jit`, `jax_inputs_batched2.shape` != `jax_inputs_batched.shape` which means that the function is compiled again. To remedy this you should make sure that all your inputs have the exact same shape, *e.g.* by padding the input. This way you won't trigger a new recompilation.<|||||>Hi @patrickvonplaten,

Thank you for the explanation. The different shapes indeed resulted in recompilation. It is now about 15% faster than PyTorch :)

Thanks for the help!