Cannot pickle linen Modules

I am using 0.3.4 and I am getting an error when trying to pickle flax modules, specifically Dense seems to be the problem but other might have similar issues.

Problem you have encountered:

from flax import linen
import pickle

with open("model.pkl", "wb") as f:
    model = pickle.dump(linen.Dense(10), f)

Traceback (most recent call last): File “test.py”, line 8, in <module> model = pickle.dump(linen.Dense(10), f) AttributeError: Can’t pickle local object ‘variance_scaling.<locals>.init’

While the previous is solved with cloudpickle, this other code doesn’t work:

import cloudpickle
from flax import linen
import pickle

class IndentityFlax(linen.Module):
    def __call__(self, x):
        return x

with open("mlp.pkl", "wb") as f:
    cloudpickle.dump(IndentityFlax(), f)

Traceback (most recent call last): File “test.py”, line 25, in <module> cloudpickle.dump(IndentityFlax(), f) File “/data/cristian/elegy/.venv/lib/python3.8/site-packages/cloudpickle/cloudpickle_fast.py”, line 55, in dump CloudPickler( File “/data/cristian/elegy/.venv/lib/python3.8/site-packages/cloudpickle/cloudpickle_fast.py”, line 563, in dump return Pickler.dump(self, obj) TypeError: cannot pickle ‘_thread._local’ object

[jheek]

@cgarciae I found the issue. cloudpickle will not serialize functions that are part of a library but it does serialize other globals. I guess this is a python limitation (missing qualname I suppose?). We use the threadlocal inside a decorator function which will get serialized. All we have to do is factor out the body of the decorator into a method so cloudpickle doens’t serialize its closure variables. I’m working on a PR

[jheek]

+1 to the comments by @levskaya

Especially remote code execution is really overlooked. This is of course intended for RPC protocols but we really don’t want a protocol with remote code execution as the standard way of distributing Flax models. I’m also curious how the PyTorch worlds thinks about this. The remote code execution problems is not mentioned in the torch.save/load docs.