Repository: google/flax

Issue:
flax.errors.ScopeParamNotFoundError: No parameter named "kernel" exists in "/Conv_0".
`Model.apply({'params':params}, batch)` in the loss function seems to throw the error above. I pretty much followed the examples in the docs line-by-line with no luck.

Here is a minimal example of the issue reproduced in google colab - https://colab.research.google.com/drive/12mRim_N4cWmv4nmeuknq8RT2VWUA5egB


Discussion:
you wrote

```
parameters = SimpleCNN6Layer(n=16).init({'params':jax.random.PRNGKey(0)}, jax.numpy.ones((16, 4000, 1)))
optimizer = optim.Adam(learning_rate=3e-4).create(parameters)
```

but you probably meant

```
variables = SimpleCNN6Layer(n=16).init({'params':jax.random.PRNGKey(0)}, jax.numpy.ones((16, 4000, 1)))
optimizer = optim.Adam(learning_rate=3e-4).create(variables['params'])
```
You have to make sure you take the `params` of the variable dict returns by `init`. So when you create your optimizer you should do:

```
variables = SimpleCNN6Layer(n=16).init({'params':jax.random.PRNGKey(0)}, jax.numpy.ones((16, 4000, 1)))
optimizer = optim.Adam(learning_rate=3e-4).create(variables['params'])
```

The error message is not very clear, so thank you for bringing this up! 👍 

We should consider improving the error message to something like "Maybe you passing in an incorrect variable dict"?
Wow, @andsteing  beat me to it with an identical code snippet 😄 
hah you clearly spent that 1 minute on a much better explanation !

👍  👍  👍  

Resolving patch:
diff --git a/flax/errors.py b/flax/errors.py
--- a/flax/errors.py
+++ b/flax/errors.py
@@ -141,9 +141,10 @@ class ApplyScopeInvalidVariablesStructureError(FlaxError):
   For more explanation on variable dicts, please see :mod:`flax.core.variables`.
   """
   def __init__(self, variables):
-    super().__init__('Expected the first argument passed to an apply function '
-                     'to be a dictionary containing a \'params\' key at the '
-                     f'root level, but got "{variables}".')
+    super().__init__('Expect the `variables` (first argument) passed to apply() '
+                     'to be a dict with the structure {"params": ...}, but got a dict '
+                     'with an extra params layer, i.e.  {"params": {"params": ... } }. '
+                     f'You should instead pass in your dict\'s ["params"].')
 
 
 class ScopeParamNotFoundError(FlaxError):
@@ -160,16 +161,18 @@ class Embed(nn.Module):
       def __call__(self, inputs, embed_name='embedding'):
         inputs = inputs.astype('int32')
         embedding = self.param(embed_name,
-                               lecun_normal(),
+                               jax.nn.initializers.lecun_normal(),
                                (self.num_embeddings, self.features))
         return embedding[inputs]
 
-    variables = Embed(4, 8).init(random.PRNGKey(0), jnp.ones((5, 5, 1)))
-    _ = Embed().apply(variables, jnp.ones((5, 5, 1)), 'embed')
+    model = Embed(4, 8)
+    variables = model.init(random.PRNGKey(0), jnp.ones((5, 5, 1)))
+    _ = model.apply(variables, jnp.ones((5, 5, 1)), 'embed')
   """
   def __init__(self, param_name, scope_path):
-    super().__init__(f'No parameter named "{param_name}" exists in '
-                     f'"{scope_path}".')
+    super().__init__(
+        f'Could not find parameter named "{param_name}" in scope '
+        f'"{scope_path}".')
 
 
 class ScopeCollectionNotFound(FlaxError):
diff --git a/flax/linen/linear.py b/flax/linen/linear.py
--- a/flax/linen/linear.py
+++ b/flax/linen/linear.py
@@ -282,8 +282,9 @@ def __call__(self, inputs: Array) -> Array:
     inputs = jnp.asarray(inputs, self.dtype)
 
     if isinstance(self.kernel_size, int):
-      raise TypeError('The kernel size must be specified as a'
-                      ' tuple/list of integers (eg.: [3, 3]).')
+      raise TypeError('Expected Conv kernel_size to be a'
+                      ' tuple/list of integers (eg.: [3, 3]) but got'
+                      f' {self.kernel_size}.')
     else:
       kernel_size = tuple(self.kernel_size)
 


Test patch:
diff --git a/tests/core/core_lift_test.py b/tests/core/core_lift_test.py
--- a/tests/core/core_lift_test.py
+++ b/tests/core/core_lift_test.py
@@ -47,7 +47,7 @@ def f(scope):
                         split_rngs={'params': True})
       dense(scope.push('dense'), np.ones((3, 2)), 2)
 
-    msg = r'No parameter named "kernel" exists in "/vmap\(dense\)".'
+    msg = r'Could not find parameter named "kernel" in scope "/vmap\(dense\)".'
     with self.assertRaisesRegex(errors.ScopeParamNotFoundError, msg):
       apply(f)({'params': {'dense': {'abc': np.ones((3, 3))}}})
 
diff --git a/tests/core/core_scope_test.py b/tests/core/core_scope_test.py
--- a/tests/core/core_scope_test.py
+++ b/tests/core/core_scope_test.py
@@ -121,11 +121,11 @@ def f(scope):
         },
     })
     apply(f)(params)  # Valid.
-    msg = 'dictionary containing a \'params\' key at the root level'
+    msg = 'but got a dict with an extra params layer'
     with self.assertRaisesRegex(errors.ApplyScopeInvalidVariablesStructureError,
                                 msg):
       apply(f)({'params': params})
-      
+
   def test_mutate_undefined_collection(self):
     def f(scope):
       scope.put_variable('state', 'test', 123)
@@ -138,7 +138,7 @@ def test_undefined_param(self):
     def f(scope):
       nn.dense(scope.push('dense'), np.ones((1, 2)), 2)
 
-    msg = r'No parameter named "kernel" exists in "/dense".'
+    msg = r'Could not find parameter named "kernel" in scope "/dense".'
     with self.assertRaisesRegex(errors.ScopeParamNotFoundError, msg):
       apply(f)({'params': {'abc': 1}})