IPython autoreload is not supported by flax

[n-gao]

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.

%load_ext autoreload
%autoreload 2
<load file>
<edit file>
<call any new cell>

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:

[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__=<class 'ferminet.FeatureConstruction'> 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)
]