Flax linnen model init not working on parallelised code

Provide as much information as possible. At least, this should include a description of your issue and steps to reproduce the problem. If possible also provide a summary of what steps or workarounds you have already tried.

### System information
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Terminal in JupyterLab

- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`:
Name: flax
Version: 0.5.3
Summary: Flax: A neural network library for JAX designed for flexibility
Home-page: https://github.com/google/flax
Author: Flax team
Author-email: flax-dev@google.com
License: UNKNOWN
Location: /home/fernanda/.local/lib/python3.8/site-packages
Requires: jax, matplotlib, msgpack, numpy, optax, PyYAML, rich, tensorstore, typing-extensions
Required-by: orbax
---
Name: jax
Version: 0.4.6
Summary: Differentiate, compile, and transform Numpy code.
Home-page: https://github.com/google/jax
Author: JAX team
Author-email: jax-dev@google.com
License: Apache-2.0
Location: /home/fernanda/.local/lib/python3.8/site-packages
Requires: numpy, opt-einsum, scipy
Required-by: chex, diffrax, equinox, flax, optax, orbax, orbax-checkpoint, richmol
---
Name: jaxlib
Version: 0.4.6
Summary: XLA library for JAX
Home-page: https://github.com/google/jax
Author: JAX team
Author-email: jax-dev@google.com
License: Apache-2.0
Location: /home/fernanda/.local/lib/python3.8/site-packages
Requires: numpy, scipy
Required-by: chex, optax, orbax, orbax-checkpoint

- Python version: 3.8

### Problem you have encountered:

When using paralellisation with joblib Parallel, delayed, recently the code is unavailable to init the parameters of the NN model. Depending on the backend that I use it will give a bug error (using loky) or it will just get stuck ('multiprocessing').

This behaviour was not previously observed.

### What you expected to happen:

Being able to initialise parameters as regularly.