From 2cc33073afbd4fdc77d88bdcd7f2f915e9b5cd83 Mon Sep 17 00:00:00 2001 From: Neil Girdhar Date: Fri, 11 Feb 2022 03:05:56 -0500 Subject: [PATCH] Support string annotations for module name When a file is marked "from __future__ import annotations", which is default in Python 3.11, annotations are strings. The check for the module name being str needs to check against the string 'str' in addition to the symbol str. --- flax/linen/module.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flax/linen/module.py b/flax/linen/module.py index 87992d361..3c914b67d 100644 --- a/flax/linen/module.py +++ b/flax/linen/module.py @@ -535,7 +535,7 @@ def _customized_dataclass_transform(cls): if ('parent' in annotations and annotations['parent'] != parent_annotation): raise errors.ReservedModuleAttributeError(annotations) - if 'name' in annotations and annotations['name'] != str: + if 'name' in annotations and annotations['name'] not in ('str', str): raise errors.ReservedModuleAttributeError(annotations) # Add `parent` and `name` default fields at end. # We temporarily modify base class __dataclass_fields__ to force desired