Merge pull request #4634 from IvyZX:parentattr

PiperOrigin-RevId: 738944800

Author: Flax Authors

flax/nnx/bridge/module.py

@@ -143,37 +143,34 @@ def _module_meta_call(cls: type[M], *args, **kwargs) -> M:
 
   name = None
   if parent_ctx is not None:
-    want_name = 'name' in kwargs and 'name' in inspect.get_annotations(cls)
-    if not want_name and not parent_ctx.in_compact and 'name' in kwargs:
-      raise ValueError(
-        f"'name' can only be set in @compact functions. If in setup(), "
-          "use parent's `self.<attr_name> to set the submodule name.")
+    if 'parent' in kwargs:
+      parent = kwargs.pop('parent')
+      if parent_ctx.in_compact and parent is not None:
+        raise ValueError(
+          f"'parent' can only be set to None, got {type(parent).__name__}"
+        )
+    else:
+      parent = parent_ctx.module
 
-    if parent_ctx.in_compact:
-      if 'parent' in kwargs:
-        parent = kwargs.pop('parent')
-        if parent is not None:
-          raise ValueError(
-            f"'parent' can only be set to None, got {type(parent).__name__}"
-          )
-      else:
-        if 'name' in kwargs:
-          name = kwargs['name'] if want_name else kwargs.pop('name')
-          if not isinstance(name, str):
-            raise ValueError(f"'name' must be a 'str', got {type(name).__name__}")
-        else:
-          name = _auto_submodule_name(parent_ctx, cls)
-        parent = parent_ctx.module
+    if 'name' in kwargs:
+      name = kwargs['name']
+      if not 'name' in inspect.get_annotations(cls):
+         kwargs.pop('name')
+      if not isinstance(name, str):
+        raise ValueError(f"'name' must be a 'str', got {type(name).__name__}")
+    elif parent_ctx.in_compact:
+      name = _auto_submodule_name(parent_ctx, cls)
 
   module = nnx_module.ModuleMeta.__call__(cls, *args, **kwargs)
   module.scope = None
   module.attr_priorities = {}
 
-  # compact behavior
   if parent is not None:
     assert parent.scope is not None
-    assert name is not None
-    setattr(parent, name, module)
+    # compact, or setup if `name` exists
+    if name is not None:
+      setattr(parent, name, module)
+      parent.set_attr_priority(name, AttrPriority.INIT_PARENT)
 
   return module  # type: ignore
 
@@ -182,9 +179,10 @@ def _module_meta_call(cls: type[M], *args, **kwargs) -> M:
 ModuleMeta.__call__ = _module_meta_call  # type: ignore
 
 class AttrPriority(enum.IntEnum):
-  HIGH = 1
-  DEFAULT = 2
-  LOW = 3
+  HIGH = 0
+  INIT_PARENT = 20
+  DEFAULT = 50
+  LOW = 100
 
 
 class PriorityStr(str):

tests/nnx/bridge/module_test.py

@@ -521,6 +521,26 @@ def __call__(self, x):
     params = model.init(jax.random.key(0), x)['params']
     self.assertSameElements(['zzz'], params.keys())
 
+  def test_linen_layer_naming(self):
+    class Dense(bridge.Module):
+      dout: int
+      @bridge.compact
+      def __call__(self, x):
+        return x @ self.param('w', lambda _: jnp.ones((x.shape[-1], self.dout)))
+
+    class MLP(bridge.Module):
+      nlayers: int
+      def setup(self):
+        self.layers = [Dense(4, name=f'layer_{i}') for i in range(self.nlayers)]
+      def __call__(self, x):
+        for layer in self.layers:
+          x = layer(x)
+        return x
+
+    model = MLP(nlayers=3)
+    x = jnp.ones((2, 4))
+    params = model.init(jax.random.key(0), x)['params']
+    self.assertSameElements([f'layer_{i}' for i in range(3)], params.keys())
 
 
 if __name__ == '__main__':