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 24, 2024
1 parent b2277ab commit c15bfe0
Showing 1 changed file with 3 additions and 0 deletions.
3 changes: 3 additions & 0 deletions flax/linen/module.py
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,9 @@ 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)
if not hasattr(cls, '_FIELDS'):
# Edge case with no fields e.g. module = nn.Module() causes error later.
return object.__repr__(module)
cls_name = cls.__name__
rep = ''

Expand Down

0 comments on commit c15bfe0

Please sign in to comment.