From ad9b06d40465b96b2dec19acaa1eb97eb97c6bca Mon Sep 17 00:00:00 2001 From: Flax Team Date: Tue, 24 Sep 2024 00:16:57 -0700 Subject: [PATCH] Fix error when printing module with no fields. --- Currently, the edge case code will cause an error: ``` print(nn.Module()) ``` PiperOrigin-RevId: 678109450 --- flax/linen/module.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/flax/linen/module.py b/flax/linen/module.py index 406912d22c..50fef7a978 100644 --- a/flax/linen/module.py +++ b/flax/linen/module.py @@ -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 = ''