Raise error if user uses 'name' in bridge module setup

flax/nnx/bridge/module.py

@@ -123,7 +123,7 @@ def _module_meta_call(cls: type[M], *args, **kwargs) -> M:
   parent = None
   module: M
 
-  if parent_ctx is not None and parent_ctx.in_compact:
+  if parent_ctx is not None:
     if 'parent' in kwargs:
       parent = kwargs.pop('parent')
       if parent is not None:
@@ -136,6 +136,10 @@ def _module_meta_call(cls: type[M], *args, **kwargs) -> M:
       parent_ctx.type_counter[cls] += 1
 
       if 'name' in kwargs:
+        if not parent_ctx.in_compact:
+          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.")
         name = kwargs.pop('name')
         if not isinstance(name, str):
           raise ValueError(f"'name' must be a 'str', got {type(name).__name__}")
@@ -202,7 +206,6 @@ def param(  # type: ignore[invalid-annotation]
     name: str,
     init_fn: tp.Callable[..., A],
     *init_args,
-    unbox: bool = True,
     **init_kwargs,
   ) -> variablelib.Param[A]:
     # TODO(cgarciae): implement same condition as linen
@@ -253,7 +256,6 @@ def variable(  # type: ignore[invalid-annotation]
     name: str,
     init_fn: tp.Callable[..., A] | None = None,
     *init_args,
-    unbox: bool = True,
     **init_kwargs,
   ) -> variablelib.Variable[A]:
     variable_type = variablelib.variable_type_from_name(

tests/nnx/bridge/wrappers_test.py

@@ -640,6 +640,14 @@ def __call__(self, x):
     y = bar.apply(variables, x)
     self.assertEqual(y.shape, (1, 5))
 
+    with self.assertRaises(ValueError):
+      class SetupBar(bridge.Module):
+        def setup(self):
+          self.xyz = Foo(5, name='xyz')
+        def __call__(self, x):
+          return self.xyz(x)
+      SetupBar().init(0, x)
+
   def test_dense_port(self):
     class Dense(bridge.Module):
       features: int