Repository: google/flax

Issue:
Fail to cloudpickle nn.Module when property is set: "maximum recursion depth exceeded"
[Colab link](https://colab.research.google.com/drive/1dBCU5vBAxlKvzsBffvIL6JvwTRfEv-RD#scrollTo=vIcLbxV_pveJ)

I'm trying to use cloudpickle to dump and reload a flax.linen.Module object, when the object has "@property" set, it would fail with `RecursionError: maximum recursion depth exceeded`.

This is the minimal code to repro:
```python
import pickle
from cloudpickle import cloudpickle_fast
from flax import linen as nn

class NNModuleWithProperty(nn.Module):
  @property
  def my_property(self):
    print('Calling my_property')
    return "my_property"

def dump_and_reload(pickle_cls, obj, filename='/tmp/test_file.pkl'):
  with open(filename, 'wb') as f:
    pickle_cls.dump(obj, f)

  with open(filename, 'rb') as f:
    obj_loaded = pickle_cls.load(f)
  return obj_loaded

dump_and_reload(cloudpickle_fast, NNModuleWithProperty())
```

It would fail with the following stack trace:
```
---------------------------------------------------------------------------
RecursionError                            Traceback (most recent call last)
[<ipython-input-6-2df3a34fadbf>](https://localhost:8080/#) in <cell line: 1>()
----> 1 dump_and_reload(cloudpickle_fast, NNModuleWithProperty())

2 frames
[<ipython-input-3-c689b3a91da9>](https://localhost:8080/#) in dump_and_reload(pickle_cls, obj, filename)
     19 
     20   with open(filename, 'rb') as f:
---> 21     obj_loaded = pickle_cls.load(f)
     22   return obj_loaded

[/usr/local/lib/python3.10/dist-packages/flax/linen/module.py](https://localhost:8080/#) in __getattr__(self, name)
    711 
    712     def __getattr__(self, name):
--> 713       return getattr(self.wrapped, name)
    714 
    715   return _DescriptorWrapper(descriptor)

... last 1 frames repeated, from the frame below ...

[/usr/local/lib/python3.10/dist-packages/flax/linen/module.py](https://localhost:8080/#) in __getattr__(self, name)
    711 
    712     def __getattr__(self, name):
--> 713       return getattr(self.wrapped, name)
    714 
    715   return _DescriptorWrapper(descriptor)

RecursionError: maximum recursion depth exceeded
```

However, it works well with pickle:
```python
dump_and_reload(pickle, NNModuleWithProperty())
```

### System information
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Ubuntu 22.04.2 LTS
- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`:
  - flax: 0.7.2
  - jax: 0.4.14
  - jaxlib: 0.4.14+cuda11.cudnn86
- Python version: 3.10.12



Discussion:

Resolving patch:
diff --git a/flax/linen/module.py b/flax/linen/module.py
--- a/flax/linen/module.py
+++ b/flax/linen/module.py
@@ -862,6 +862,8 @@ def __set_name__(self, *args, **kwargs):
         self.wrapped.__set_name__(*args, **kwargs)
 
     def __getattr__(self, name):
+      if 'wrapped' not in vars(self):
+        raise AttributeError()
       return getattr(self.wrapped, name)
 
   return _DescriptorWrapper(descriptor)


Test patch:
diff --git a/tests/linen/linen_module_test.py b/tests/linen/linen_module_test.py
--- a/tests/linen/linen_module_test.py
+++ b/tests/linen/linen_module_test.py
@@ -22,6 +22,7 @@
 import inspect
 import operator
 import sys
+from tempfile import TemporaryDirectory
 from typing import (
     Any,
     Callable,
@@ -2399,6 +2400,31 @@ def record_interceptor(f, args, kwargs, context):
     self.assertIs(called[0], bar)
     self.assertIs(called[1], foo)
 
+  def test_cloudpickle_module(self):
+    from cloudpickle import cloudpickle_fast
+
+    class NNModuleWithProperty(nn.Module):
+      a: int
+      b: str
+
+      @property
+      def my_property(self):
+        return self.b * self.a
+
+    m = NNModuleWithProperty(a=2, b='ok')
+
+    with TemporaryDirectory() as tmpdir:
+      filename = f'{tmpdir}/module.pkl'
+      with open(filename, 'wb') as f:
+        cloudpickle_fast.dump(m, f)
+
+      with open(filename, 'rb') as f:
+        obj_loaded = cloudpickle_fast.load(f)
+
+    self.assertEqual(obj_loaded.a, 2)
+    self.assertEqual(obj_loaded.b, 'ok')
+    self.assertEqual(obj_loaded.my_property, 'okok')
+
 
 class LeakTests(absltest.TestCase):