Repository: google/flax

Issue:
Module's `__repr__` not showing parents' fields in its attributes
### System information
- `macOS 13.1`
- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`:
```
Name: flax
Version: 0.6.1
---
Name: jax
Version: 0.4.1
---
Name: jaxlib
Version: 0.4.1
```
- `Python 3.10.8`

### Problem you have encountered:
For a module inheriting other modules or dataclasses, when `__repr__` is called, it only prints its own fields, but none of its parent classes' fields, under `# attributes`. 
These fields are not the same as the signature of the generated `__init__` constructor either.


### What you expected to happen:
The `attributes` should have the same fields as the generated `__init__` constructor (without `parent` and `name`). From the module user's point of view, since I need to fill these attributes when init the module, I expect to see them all when printing the module.


### Steps to reproduce:
For the following example
```Python
from flax.linen import Module
from inspect import signature


class Base1(Module):
    p11: int
    p12: int


class Base2(Module):
    p21: int
    p22: int


class Model(Base2, Base1):
    a: int
    b: int


m = Model(1, 2, 3, 4, 5, 6)
```

If we `print(signature(m.__init__))`, we get
```
(p11: int, p12: int, p21: int, p22: int, a: int, b: int, parent: Union[Type[flax.linen.module.Module], Type[flax.core.scope.Scope], Type[flax.linen.module._Sentinel], NoneType] = <flax.linen.module._Sentinel object at 0x12b7d7700>, name: str = None) -> None
```
but for `print(m)`, we get
```
Model(
    # attributes
    a = 5
    b = 6
)
```
I think it should be
```
Model(
    # attributes
    p11 = 1
    p12 = 2
    p21 = 3
    p22 = 4
    a = 5
    b = 6
)
```




Discussion:
[A quick fix](https://github.com/google/flax/compare/main...JyChang012:flax:fix/repr_attributes) that uses `__dataclass_fields__` to find all fields
Feel free to close this if current behavior of `__repr__` is intended. @cgarciae 

Resolving patch:
diff --git a/flax/linen/module.py b/flax/linen/module.py
--- a/flax/linen/module.py
+++ b/flax/linen/module.py
@@ -97,8 +97,9 @@ def _module_repr(module: 'Module', num_spaces: int = 4):
   cls = type(module)
   cls_name = cls.__name__
   rep = ''
-  attributes = {k: v for k, v in cls.__annotations__.items()
-                if k not in ('parent', 'name')}
+
+  attributes = {f.name: f.type for f in dataclasses.fields(cls)
+                if f.name not in ('parent', 'name')}
   child_modules = {k: v for k, v in module._state.children.items()  # pytype: disable=attribute-error
                    if isinstance(v, Module)}
   if attributes:


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
@@ -1897,6 +1897,24 @@ def __call__(self):
                                 'Trying to access a property that'):
       foo.apply({})
 
+  def test_repr(self):
+
+    class Base1(nn.Module):
+      a: int
+
+    class Base2(nn.Module):
+        b: str
+
+    class Foo(Base2, Base1):
+        c: float
+
+    module = Foo(a=1, b='ok', c=3.0)
+    str_rep = repr(module)
+
+    self.assertIn('a = 1', str_rep)
+    self.assertIn("b = 'ok'", str_rep)
+    self.assertIn('c = 3.0', str_rep)
+
 
 class LeakTests(absltest.TestCase):