undo recent typing changes as it may be confusing pyright

PiperOrigin-RevId: 559463753

Author: levskaya

flax/linen/module.py

@@ -873,10 +873,6 @@ def __getattr__(self, name):
 # -----------------------------------------------------------------------------
 
 
-def module_field(*, kw_only: bool = False, default: Optional[Any] = ...) -> Any:
-  ...
-
-
 # The ModuleBase class is created only to make static analyzers happy
 # mainly pytype and pyright. Some notes:
 # * pyright (correctly) complains that Module itself is not a dataclass, even
@@ -891,12 +887,14 @@ def module_field(*, kw_only: bool = False, default: Optional[Any] = ...) -> Any:
 # * Other attributes are annotated for completeness. Because we are using
 #   the `if typing.TYPE_CHECKING` pattern, these annotations are not present
 #   at runtime so they don't affect the dataclass behavior.
-@dataclass_transform(field_specifiers=(module_field,))  # type: ignore[literal-required]
+@dataclass_transform()
 class ModuleBase:
   if typing.TYPE_CHECKING:
+    name: Optional[str] = None
     scope: Optional[Scope]
     _state: _ModuleInternalState
     _parent_ref: Union['Module', weakref.ReferenceType['Module'], None]
+    parent: Union['Module', _Sentinel, None]
     __dataclass_fields__: Dict[str, dataclasses.Field]
 
 
@@ -936,10 +934,6 @@ def __call__(self, x):
   """
 
   if typing.TYPE_CHECKING:
-    name: Optional[str] = module_field(kw_only=True, default=None)
-    parent: Union['Module', _Sentinel, None] = module_field(
-        kw_only=True, default=None
-    )
 
     def __init__(self, *args, **kwargs):
       # this stub makes sure pytype accepts constructor arguments.