Memory accumulation when using `from_bytes`

Hi all! If I load a flax model from disk multiple times using `from_bytes` (I'm using `diffusers`'s `from_pretrained` method here which uses `from_bytes` underneath) I see my RAM memory increase by quite a bit. Here's a simple reproduction using a small model (~300MB):

```python
from diffusers import FlaxAutoencoderKL

def get_memory_usage_percentage():
with open('/proc/meminfo', 'r') as file:
for line in file:
if 'MemAvailable:' in line:
mem_available = int(line.split()[1])
elif 'MemTotal:' in line:
mem_total = int(line.split()[1])
mem_used = (mem_total - mem_available)
return mem_used / 1024**2

def log_memory(run):
mem_used = get_memory_usage_percentage()
print(f'Memory usage at run "{run}": {mem_used:.2f}GB')

for run in range(10):
log_memory(run)
vae, vae_params = FlaxAutoencoderKL.from_pretrained('runwayml/stable-diffusion-v1-5', subfolder='vae', revision='flax')
```

When I run this on a TPU (v3-8) I get
```
Memory usage at run "0": 4.92GB
Memory usage at run "1": 7.36GB
Memory usage at run "2": 8.02GB
Memory usage at run "3": 8.31GB
Memory usage at run "4": 8.31GB
Memory usage at run "5": 8.31GB
Memory usage at run "6": 8.32GB
Memory usage at run "7": 8.33GB
Memory usage at run "8": 8.33GB
Memory usage at run "9": 8.33GB
```
Given the model is fairly small I don't understand why RAM usage is so high in this case? Also, I wonder why memory is not released after the second iteration and is increasing again until it plateaus.

On a CPU/[Colab Notebook](https://colab.research.google.com/drive/1diE9-W_5L1fwdmzbcdQskk8F1JCHGNWm?usp=sharing) I get the following:
```
Memory usage at run "0": 1.43GB
Memory usage at run "1": 1.64GB
Memory usage at run "2": 2.04GB
Memory usage at run "3": 2.42GB
Memory usage at run "4": 2.43GB
Memory usage at run "5": 2.44GB
Memory usage at run "6": 2.44GB
Memory usage at run "7": 2.44GB
Memory usage at run "8": 2.43GB
Memory usage at run "9": 2.44GB
```
Which follows a similar pattern at a smaller scale.

I'm not sure this is a bug - would be great to better understand what's going on. I've had similar odd memory behavior with `jax.device_get` recently ([see here](https://github.com/google/jax/discussions/15972)), which may be related?

### System information
- Ubuntu 18
- `flax==0.6.9`, `jax==0.4.8`, `jaxlib==0.4.8`
- Python version: 3.10.11