Skip to content

Commit

Permalink
Fix error when printing module with no fields.
Browse files Browse the repository at this point in the history
---

Currently, the edge case code will cause an error:

```
print(nn.Module())
```

PiperOrigin-RevId: 678109450
  • Loading branch information
Flax Team committed Sep 27, 2024
1 parent 37f643d commit db129cc
Showing 1 changed file with 6 additions and 1 deletion.
7 changes: 6 additions & 1 deletion flax/linen/module.py
Original file line number Diff line number Diff line change
Expand Up @@ -111,12 +111,17 @@ def _attr_repr(value: Any):
def _module_repr(module: 'Module', num_spaces: int = 4):
"""Returns a pretty printed representation of the module."""
cls = type(module)
try:
fields = dataclasses.fields(cls)
except TypeError:
# Edge case with no fields e.g. module = nn.Module() causes error later.
return object.__repr__(module)
cls_name = cls.__name__
rep = ''

attributes = {
f.name: f.type
for f in dataclasses.fields(cls)
for f in fields
if f.name not in ('parent', 'name') and f.repr
}
child_modules = {
Expand Down

0 comments on commit db129cc

Please sign in to comment.