Repository: google/flax

Issue:
AttributeError isn't raised correctly from properties of linen modules
### System information
- OS Platform and Distribution: ```Any```
- Flax, jax, jaxlib versions ```flax==0.6.1, jax==0.3.23, jaxlib==0.3.22```
- Python version: ```3.10```
- GPU/TPU model and memory: ```N/A```
- CUDA version (if applicable): ```N/A```

### Problem you have encountered:
Linen modules do not seem to report ```AttributeError```s correctly for non-existent attributes accessed by properties. Instead, the raised error claims that the property itself does not exist. 

This issue doesn't seem to occur in linen module methods, and to the best of my knowledge, it doesn't occur with other types of errors in linen properties, such as with mismatched einsum indices. 

### What you expected to happen:
Properties in linen modules would ideally perform error reporting the same way as properties of ordinary classes. E.g.,
```
class Foo:
    def __init__(self):
        self.bar = 0
    @property
    def prop(self):
        return self.baz
    def __call__(self):
        return self.prop

foo = Foo()
foo()
```

which gives
```
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<stdin>", line 8, in __call__
  File "<stdin>", line 6, in prop
AttributeError: 'Foo' object has no attribute 'baz'
```

### Steps to reproduce:

However, Linen modules do not perform error reporting correctly for properties: 
```
import flax.linen as nn
import jax
import jax.numpy as jnp

class Foo(nn.Module):
    def setup(self):
        self.bar = self.param("bar", jax.nn.initializers.normal(0.01), [10, 10], jnp.float32)
    @property
    def prop(self):
        return self.baz
    def __call__(self, inputs):
        return self.prop

foo = Foo()
params = foo.init(
    {"params": jax.random.PRNGKey(0)}, 
    inputs=jnp.ones(shape=[64, 128], dtype=jnp.float32)
)["params"]
```

### Logs, error messages, etc:

Running the code immediately above gives
```
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/lucaslingle/opt/miniconda3/envs/some_project/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/lucaslingle/opt/miniconda3/envs/some_project/lib/python3.10/site-packages/flax/linen/module.py", line 1333, in init
    _, v_out = self.init_with_output(
  File "/Users/lucaslingle/opt/miniconda3/envs/some_project/lib/python3.10/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/lucaslingle/opt/miniconda3/envs/some_project/lib/python3.10/site-packages/flax/linen/module.py", line 1289, in init_with_output
    return init_with_output(
  File "/Users/lucaslingle/opt/miniconda3/envs/some_project/lib/python3.10/site-packages/flax/core/scope.py", line 897, in wrapper
    return apply(fn, mutable=mutable, flags=init_flags)({}, *args, rngs=rngs,
  File "/Users/lucaslingle/opt/miniconda3/envs/some_project/lib/python3.10/site-packages/flax/core/scope.py", line 865, in wrapper
    y = fn(root, *args, **kwargs)
  File "/Users/lucaslingle/opt/miniconda3/envs/some_project/lib/python3.10/site-packages/flax/linen/module.py", line 1750, in scope_fn
    return fn(module.clone(parent=scope), *args, **kwargs)
  File "/Users/lucaslingle/opt/miniconda3/envs/some_project/lib/python3.10/site-packages/flax/linen/module.py", line 402, in wrapped_module_method
    return self._call_wrapped_method(fun, args, kwargs)
  File "/Users/lucaslingle/opt/miniconda3/envs/some_project/lib/python3.10/site-packages/flax/linen/module.py", line 705, in _call_wrapped_method
    y = fun(self, *args, **kwargs)
  File "<stdin>", line 8, in __call__
  File "/Users/lucaslingle/opt/miniconda3/envs/some_project/lib/python3.10/site-packages/flax/linen/module.py", line 783, in __getattr__
    raise AttributeError(msg)
jax._src.traceback_util.UnfilteredStackTrace: AttributeError: "Foo" object has no attribute "prop".
```


Discussion:
@lucaslingle thanks for reporting this. Its not that we are not reporting an AttributeError for the incorrect attribute, it seems that we currently don't support properties at all. I'll look into this.
Update, after some digging, it seems this error will happen with any code that has the following properties:
* Class implements `__getattr__`
* The `property` raises a `AttributeError`

We will look a way to wrap properties to catch this attribute error an launch a custom error.

Resolving patch:
diff --git a/flax/errors.py b/flax/errors.py
--- a/flax/errors.py
+++ b/flax/errors.py
@@ -629,6 +629,23 @@ def __call__(self, input):
   def __init__(self):
     super().__init__('Overrode `.__post_init__()` without calling `super().__post_init__()`')
 
+class DescriptorAttributeError(FlaxError):
+  """
+  This error occurs when you are trying to access a property that is accessing a non-existent attribute.
+  For example, the error will be raised when trying to run this code::
+
+    class Foo(nn.Module):
+        @property
+        def prop(self):
+            return self.non_existent_field # ERROR!
+        def __call__(self, x):
+            return self.prop
+
+    foo = Foo()
+    variables = foo.init(jax.random.PRNGKey(0), jnp.ones(shape=(1, 8)))
+  """
+  def __init__(self):
+    super().__init__('Trying to access a property that is accessing a non-existent attribute.')
 
 class InvalidCheckpointError(FlaxError):
   """
diff --git a/flax/linen/module.py b/flax/linen/module.py
--- a/flax/linen/module.py
+++ b/flax/linen/module.py
@@ -28,8 +28,8 @@
 import jax
 import numpy as np
 import jax.numpy as jnp
-from typing_extensions import \
-    dataclass_transform  # pytype: disable=not-supported-yet
+from typing_extensions import Protocol, \
+  dataclass_transform  # pytype: disable=not-supported-yet
 
 import flax
 from flax import (config, core, errors, serialization, traceback_util,
@@ -386,6 +386,27 @@ def _get_local_method_names(cls: Any,
         true_methods.add(m)
   return tuple(true_methods.difference(set(exclude)))
 
+def _get_local_descriptor_names(cls: Any,
+                                exclude: Iterable[str] = ()) -> Tuple[str, ...]:
+  """Gets descriptor names of a class.
+
+  Args:
+    cls: The class to get property names for.
+    exclude: Names to exclude from output.
+  Returns:
+    A list of property names.
+  """
+  true_properties = set()
+  for m, attr in cls.__dict__.items():
+    if not callable(attr) and (
+      hasattr(attr, '__get__') or hasattr(attr, '__set__') or
+      hasattr(attr, '__delete__')
+    ):
+      mtype = type(attr)
+      if mtype != staticmethod and mtype != classmethod:
+        true_properties.add(m)
+  return tuple(true_properties.difference(set(exclude)))
+
 
 def wrap_method_once(fun: Callable[..., Any]) -> Callable[..., Any]:
   """Manages Module state for a given user-defined method.
@@ -414,6 +435,20 @@ def wrapped_module_method(*args, **kwargs):
   wrapped_module_method.method_handler_wrapped = True  # type: ignore[attr-defined]
   return wrapped_module_method
 
+def wrap_descriptor_once(descriptor) -> "DescriptorWrapper":
+  """Wraps a descriptor to give better error messages.
+
+  Args:
+    prop: User-defined Module attribute descriptor.
+  Returns:
+    Wrapped descriptor.
+  """
+  # Don't rewrap descriptors.
+  if isinstance(descriptor, DescriptorWrapper):
+    return descriptor
+
+  return create_descriptor_wrapper(descriptor)
+
 
 def _wrap_hash(hash_fn: Callable[..., Any]) -> Callable[..., Any]:
   """Wraps a hash function with some check for Flax Modules."""
@@ -552,6 +587,55 @@ def __set__(self, obj, value):
     object.__setattr__(obj, "_parent_ref", maybe_weak)
 
 
+class Descriptor(Protocol):
+  __isabstractmethod__: bool
+  def __get__(self, obj, objtype=None) -> Any: ...
+  def __set__(self, obj, value) -> None: ...
+  def __delete__(self, obj) -> None: ...
+  def __set_name__(self, owner, name) -> None: ...
+
+class DescriptorWrapper:
+  pass
+
+def create_descriptor_wrapper(descriptor: Descriptor):
+  """Creates a descriptor wrapper that calls a get_fn on the descriptor."""
+
+  class _DescriptorWrapper(DescriptorWrapper):
+    """A descriptor that can wrap any descriptor"""
+
+    if hasattr(descriptor, '__isabstractmethod__'):
+      __isabstractmethod__ = descriptor.__isabstractmethod__
+
+    def __init__(self, wrapped: Descriptor):
+      self.wrapped = wrapped
+
+    # conditionally define descriptor methods
+    if hasattr(descriptor, '__get__'):
+      def __get__(self, *args, **kwargs):
+        # here we will catch internal AttributeError and re-raise it as a
+        # more informative and correct error message.
+        try:
+          return self.wrapped.__get__(*args, **kwargs)
+        except AttributeError as e:
+          raise errors.DescriptorAttributeError() from e
+
+    if hasattr(descriptor, '__set__'):
+      def __set__(self, *args, **kwargs):
+        return self.wrapped.__set__(*args, **kwargs)
+
+    if hasattr(descriptor, '__delete__'):
+      def __delete__(self, *args, **kwargs):
+        return self.wrapped.__delete__(*args, **kwargs)
+
+    if hasattr(descriptor, '__set_name__'):
+      def __set_name__(self, *args, **kwargs):
+        self.wrapped.__set_name__(*args, **kwargs)
+
+    def __getattr__(self, name):
+      return getattr(self.wrapped, name)
+
+  return _DescriptorWrapper(descriptor)
+
 # Base Module definition.
 # -----------------------------------------------------------------------------
 
@@ -613,7 +697,7 @@ def __init_subclass__(cls, **kwargs: Any) -> None:
     # We wrap user-defined methods including setup and __call__ to enforce
     # a number of different checks and to provide clear error messages.
     cls._verify_single_or_no_compact()
-    cls._wrap_module_methods()
+    cls._wrap_module_attributes()
     # Set empty class defaults.
     cls._state = _uninitialized_module_internal_state # type: ignore[attr-defined]
     cls.scope: Optional[Scope] = None # type: ignore
@@ -673,16 +757,29 @@ def _verify_single_or_no_compact(cls):
       raise errors.MultipleMethodsCompactError()
 
   @classmethod
-  def _wrap_module_methods(cls):
-    """Wraps user-defined non-inherited methods with state management functions."""
-    exclusions = ([f.name for f in dataclasses.fields(cls)] +
+  def _wrap_module_attributes(cls):
+    """Wraps user-defined non-inherited methods and descriptors with state
+    management functions.
+    """
+    # wrap methods
+    method_exclusions = ([f.name for f in dataclasses.fields(cls)] +
                   ['__eq__', '__repr__', '__init__', '__hash__',
                    '__post_init__'])
-    for key in _get_local_method_names(cls, exclude=exclusions):
+    for key in _get_local_method_names(cls, exclude=method_exclusions):
       method = getattr(cls, key)
       if hasattr(method, 'nowrap'):
         continue
       setattr(cls, key, wrap_method_once(method))
+
+    # wrap descriptors
+    descriptor_exclusions = ([f.name for f in dataclasses.fields(cls)] +
+                             ['parent', '__dict__'])
+    for key in _get_local_descriptor_names(cls, descriptor_exclusions):
+      # don't use getattr here, since it will call the descriptor
+      descriptor = cls.__dict__[key]
+      if hasattr(descriptor, 'nowrap'):
+        continue
+      setattr(cls, key, wrap_descriptor_once(descriptor))
     return cls
 
   def _call_wrapped_method(self, fun, args, kwargs):


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
@@ -1809,6 +1809,50 @@ class Network(nn.Module):
     self.assertEqual(type_hints['layers'], int)
 
 
+  def test_incorrect_property(self):
+    class Foo(nn.Module):
+      @property
+      def prop(self):
+        return self.non_existent
+      def __call__(self):
+        return self.prop
+
+    foo = Foo()
+    with self.assertRaisesRegex(
+      errors.DescriptorAttributeError, 'Trying to access a property that'):
+      foo.apply({})
+
+  def test_custom_descriptor(self):
+    class Descriptor:
+      def __get__(self, obj, objtype=None):
+        return 10
+
+    class Foo(nn.Module):
+      prop = Descriptor()
+
+      def __call__(self):
+        return self.prop
+
+    foo = Foo()
+    res = foo.apply({})
+    self.assertEqual(res, 10)
+
+  def test_custom_descriptor_error(self):
+    class Descriptor:
+      def __get__(self, obj, objtype=None):
+        return obj.non_existent
+
+    class Foo(nn.Module):
+      prop = Descriptor()
+
+      def __call__(self):
+        return self.prop
+
+    foo = Foo()
+    with self.assertRaisesRegex(
+      errors.DescriptorAttributeError, 'Trying to access a property that'):
+      foo.apply({})
+
 class LeakTests(absltest.TestCase):
 
   def test_tracer_leaks(self):