Repository: google/flax

Issue:
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:
```python
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:
```python
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



Discussion:
Ok, this exact example seems to work with `cloudpickle` but I am getting another error serializing a `flax.linen.Module` object. I will try to get a reproducible example.

So I found a minimal example that doesn't work with `cloudpickle`which seems to be what is affecting me on my actual problem, see updated issue.
(cloud)pickle issues are a little tricky. For some reason cloudpickle tries to serialize Flax internals. I spend some time before looking into it but my main issue with cloudpickle is that there doesn't seem to be a specification of their algorithm and of course the implementation is black-magic Python. I think the minimal thing we need to officially support a library like cloudpickle is a guide that explains what constraints we should adhere to in order to support cloudpickle. Perhaps something like this does exist but I couldn't find anything last time I looked for it.

You could of course also raise an issue with the cloudpickle team to see if this is even expected behavior from their side in the first place.
@jheek You happen to know which Flax internal object its trying to serialize? I am a bit hesitant to ping cloudpickle without a reproducible example that doesn't involve a whole library (flax) as part of it.

If flax users are not using (cloud)pickle, what is the current recommended way to serialize flax models? 
Yeah I agree we should try to minimize the repro. I tried out your pickle example and I was able to remove Flax from the equation:
```
init_fn = jax.nn.initializers.lecun_normal()
with open("model.pkl", "wb") as f:
    model = pickle.dump(init_fn, f)
```
So here it's really JAX that is producing a partial function that cannot be pickled.

For cloudpickle I'm not so sure what's going on but it essentially finds an internal ThreadLocal object and decides that it needs to serialize it. This I think doesn't make sense. After all only library methods touch this object (which it shouldn't serialize) and the ThreadLocal object is itself is defined top-level in the module so again it shouldn't try to serialize this object. This pattern of having state in a ThreadLocal object is quite common in Python so I think this should really be fixed in cloudpickle but perhaps I'm overlooking some edge case in how we implemented this in Flax.

@jheek thanks for looking into this!

Its very weird as you say since both `elegy` and `haiku` Modules use `ThreadLocal` but only `flax` is having issues with `cloudpickle`. I am more interested about `cloudpickle` than `pickle` since its generally more robust and pickle doesn't work for `haiku` and `elegy` either so its not really an option.

I will send a PR to `flax` with a test using `cloudpickle` to make this effort a little more formal and maybe others can try to give it a shot if that is OK with the flax team.
I am curious indeed why other libraries that use `ThreadLocal` don't have this problem...
@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

Resolving patch:
diff --git a/flax/linen/module.py b/flax/linen/module.py
--- a/flax/linen/module.py
+++ b/flax/linen/module.py
@@ -262,39 +262,9 @@ def wrapped_module_method(*args, **kwargs):
     # otherwise call the wrapped function as is.
     if args and isinstance(args[0], Module):
       self, args = args[0], args[1:]
+      return self._call_wrapped_method(fun, args, kwargs)
     else:
       return fun(*args, **kwargs)
-    is_compact_method = hasattr(fun, 'compact')
-    is_setup_method = fun.__name__ == 'setup'
-    # We lazily call setup() only when needed.
-    if is_setup_method:
-      is_recurrent = self._state.in_setup
-      self._state.in_setup = True
-    else:
-      self._try_setup()
-
-    if is_compact_method:
-      if self.scope is None:
-        raise errors.CallCompactUnboundModuleError()
-      is_recurrent = self._state.in_compact_method
-      self._state.in_compact_method = True
-    _context.module_stack.append(self)
-    try:
-      y = fun(self, *args, **kwargs)
-      if _context.capture_stack:
-        filter_fn = _context.capture_stack[-1]
-        if filter_fn and filter_fn(self, fun.__name__):
-          self.sow('intermediates', fun.__name__, y)
-      return y
-    finally:
-      _context.module_stack.pop()
-      if is_compact_method:
-        object.__setattr__(self, 'scope', self.scope.rewound())
-      # setup or compact calls can be recurrent for example due to super calls
-      # resetting the state would cause is compact/setup method
-      # to be set to False prematurely.
-      if (is_compact_method or is_setup_method) and not is_recurrent:
-        self._state.reset()
   wrapped_module_method.method_handler_wrapped = True
   return wrapped_module_method
 
@@ -523,6 +493,46 @@ def _wrap_module_methods(cls):
       setattr(cls, key, wrapped_method)
     return cls
 
+  def _call_wrapped_method(self, fun, args, kwargs):
+    """"Calls a wrapped method.
+    
+    This function is responsible for setting up the thread local state
+    correctly before calling the method and cleaning up afterwards.
+    This includes storing intermediates, setup of the compact scope,
+    and making sure setup is called before any other method.
+    """
+    is_compact_method = hasattr(fun, 'compact')
+    is_setup_method = fun.__name__ == 'setup'
+    # We lazily call setup() only when needed.
+    if is_setup_method:
+      is_recurrent = self._state.in_setup
+      self._state.in_setup = True
+    else:
+      self._try_setup()
+
+    if is_compact_method:
+      if self.scope is None:
+        raise errors.CallCompactUnboundModuleError()
+      is_recurrent = self._state.in_compact_method
+      self._state.in_compact_method = True
+    _context.module_stack.append(self)
+    try:
+      y = fun(self, *args, **kwargs)
+      if _context.capture_stack:
+        filter_fn = _context.capture_stack[-1]
+        if filter_fn and filter_fn(self, fun.__name__):
+          self.sow('intermediates', fun.__name__, y)
+      return y
+    finally:
+      _context.module_stack.pop()
+      if is_compact_method:
+        object.__setattr__(self, 'scope', self.scope.rewound())
+      # setup or compact calls can be recurrent for example due to super calls
+      # resetting the state would cause is compact/setup method
+      # to be set to False prematurely.
+      if (is_compact_method or is_setup_method) and not is_recurrent:
+        self._state.reset()
+
   def __setattr__(self, name: str, val: Any):
     """Sets an attribute on this Module.
 


Test patch: