IPython autoreload is not supported by flax

### System information
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 20.04
- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`: flax 0.4, jax 0.3.1, jaxlib 0.30 (cuda11)
- Python version: 3.8
- GPU/TPU model and memory: GTX 1080ti
- CUDA version (if applicable): 11

### Problem you have encountered:
If you want to prototype in a file and use the `autoreload` package of `ipython`, anytime you edit a file with a flax linen module you get an error.
```python
%load_ext autoreload
%autoreload 2

```
A workaround would be to copy the file content to a jupyter notebook cell and copy it back when the edits have been done.

### What you expected to happen:
The expected behavior would be that either the module class changes. I guess this behavior conflicts with linen foundations, maybe it is easier to implement it such that the newly created instances adopt the updated implementation?

### Logs, error messages, etc:
```python
[autoreload of mabenet.nn.ferminet failed: Traceback (most recent call last):
File "/nfs/staff-ssd/gaoni/miniconda3/envs/pesnet/lib/python3.8/site-packages/IPython/extensions/autoreload.py", line 257, in check
superreload(m, reload, self.old_objects)
File "/nfs/staff-ssd/gaoni/miniconda3/envs/pesnet/lib/python3.8/site-packages/IPython/extensions/autoreload.py", line 480, in superreload
update_generic(old_obj, new_obj)
File "/nfs/staff-ssd/gaoni/miniconda3/envs/pesnet/lib/python3.8/site-packages/IPython/extensions/autoreload.py", line 377, in update_generic
update(a, b)
File "/nfs/staff-ssd/gaoni/miniconda3/envs/pesnet/lib/python3.8/site-packages/IPython/extensions/autoreload.py", line 345, in update_class
update_instances(old, new)
File "/nfs/staff-ssd/gaoni/miniconda3/envs/pesnet/lib/python3.8/site-packages/IPython/extensions/autoreload.py", line 303, in update_instances
ref.__class__ = new
File "/nfs/staff-ssd/gaoni/miniconda3/envs/pesnet/lib/python3.8/site-packages/flax/linen/module.py", line 669, in __setattr__
raise errors.SetAttributeFrozenModuleError(self.__class__.__name__, name,
flax.errors.SetAttributeFrozenModuleError: Can't set __class__= for Module of type FeatureConstruction: Module instance is frozen outside of setup method. (https://flax.readthedocs.io/en/latest/flax.errors.html#flax.errors.SetAttributeFrozenModuleError)
]
```