FlaxCLIPModel memory leak due to JAX `jit` function cache
Author: josephrocca
Pull request: no

## Environment info

- `transformers` version: 4.10.2
- Platform: Linux-5.11.0-7620-generic-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.6.9
- PyTorch version (GPU?): 1.9.0+cu102 (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: FlaxCLIPModel

The problem arises when using an altered version of the official example. I'm using the outputs of `CLIPProcessor` in a `score` function that outputs how similar the text and image are (basically the same as the official example script). The difference is that I'm `jax.jit`ing the `score` function in a `for` loop and executing it with a different-sized image each time. This causes a memory leak as the previously cached `jit`s are for some reason not being cleared.

Why `jit` the `score` function in a loop? Why not just `jit` once at the start and let JAX automatically re-`jit` based on the input array shape? That unfortunately doesn't fix the issue. That's how I had the code originally, but the memory still grows. I hoped that by completely dropping the reference to the function and then re-`jit`ing it whenever I needed to change to a new image size, it would clear the old `jit` cache. That turns out not to be the case.

Here's the code that reproduces this with some memory logging stripped out for ease of reading. See the Gist/Colab links below for the full code. All I'm doing here is repeatedly re-`jit`ing the `score` function for different image sizes. This results in the script consuming more and more GPU memory until it crashes (after about 6 rounds on my machine). This example code is based on real code where I need to change the size of the images that are passed to `score` every now and then.

```python
import os
os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
os.environ["XLA_PYTHON_CLIENT_ALLOCATOR"] = "platform"

!pip install --upgrade pip
!pip install --upgrade "jax[cuda111]" flax==0.3.4 transformers==4.10.2 Pillow==8.3.2 numpy==1.19.5 ftfy==6.0.3 -f https://storage.googleapis.com/jax-releases/jax_releases.html

import jax
from jax import jit, grad
import jax.numpy as np
from jax.interpreters import xla
import PIL
from PIL import Image
import ftfy
import time
from transformers import CLIPProcessor, FlaxCLIPModel

model = FlaxCLIPModel.from_pretrained("openai/clip-vit-base-patch32")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

def score(pixel_values, input_ids, attention_mask):
    pixel_values = jax.image.resize(pixel_values, (1, 3, 224, 224), "nearest")
    inputs = {"pixel_values":pixel_values, "input_ids":input_ids, "attention_mask":attention_mask}
    outputs = model(**inputs)
    return outputs.logits_per_image[0][0]

size = 224
image = Image.new('RGB', (size, size), color=(0, 0, 66))
data = processor(text=["blue"], images=[image], return_tensors="jax", padding=True)
pixel_values = data.pixel_values
input_ids = data.input_ids
attention_mask = data.attention_mask

for i in range(100):
    print(f"\nROUND #{i}")
    
    # Decrease image size (to bust JAX jit cache):
    print("SIZE:", size)
    size -= 1
    pixel_values = jax.image.resize(pixel_values, (1, 3, size, size), "nearest")
    
    # Clear JAX GPU memory:
    jit_score = None
    time.sleep(2)
    xla._xla_callable.cache_clear()
    
    # JIT score function and execute it:
    jit_score = jit(score)   
    print(jit_score(pixel_values, input_ids, attention_mask))
```

As you can see I've tried using `XLA_PYTHON_CLIENT_ALLOCATOR=platform` as suggested by a JAX dev [here](https://github.com/google/jax/issues/1222). I've also tried `XLA_PYTHON_CLIENT_PREALLOCATE=true` (the default) and the problem remains. I'm also calling `xla._xla_callable.cache_clear()`, but that shouldn't be necessary, as mentioned by another JAX dev [here](https://github.com/google/jax/issues/2072).

## To reproduce

Here's a minimal reproduction with GPU memory logging:
* https://gist.github.com/josephrocca/d0ddd720fd4bf178337fe8d878a559d9
* https://colab.research.google.com/gist/josephrocca/d0ddd720fd4bf178337fe8d878a559d9

As you can see from the logs in the Gist, the memory accumulates/leaks every round. If I use the same `pixel_values` array, then there is no leak (predictably, since there's no re-`jit`).

I thought it might have been a problem with JAX, but I created a minimal example that's analogous to that one (but using a different and very simple `score` function instead of the Flax model) and there's no leak:
* https://gist.github.com/josephrocca/3ac6e48d81b3dc9c67b54e2f5bd2fa70
* https://colab.research.google.com/gist/josephrocca/3ac6e48d81b3dc9c67b54e2f5bd2fa70

Still seems like this might be an upstream problem - but strange that it doesn't occur with the pure JAX example. I don't know enough about the internals of FlaxCLIPModel and Flax to know whether they could be interacting with `jax.jit` in a way that's causing this 🤔

## Expected result

I'd have expected the `jit` compilation cache for a function to be cleared when the reference to that function is dropped. That seems to be the behavior of the minimal non-Flax Colab/Gist that I've linked. If all previously-cached `jit`s are kept in memory forever, then it makes it hard to use this module in non-trivial code.

Comments:
It just occurred to me that perhaps the cache *is* growing in the simplified no-Flax code example, but it's too slow to notice due to the simplicity of the function. I've just ran another test and that does indeed seem to be the case. So this is an upstream (JAX) problem or misunderstanding on my part. Here's the example showing the leak (it "leaks" about 1mb per 100 iterations):

* https://gist.github.com/josephrocca/b56f0cb99b8ac07f15bee467842788f2
* https://colab.research.google.com/gist/josephrocca/b56f0cb99b8ac07f15bee467842788f2

Closing since it doesn't seem to be a problem that's specifically caused by `transformers`.

**Edit:** I've asked how to clear the jit cache for a function in the JAX discussion forum [here](https://github.com/google/jax/discussions/7882).