Repository: google/flax

Issue:
flax.errors.ScopeParamNotFoundError: No parameter named "kernel" exists in "/MLP_0/Dense_0" when attempting to use Jax2TF with a pre-trained JAX NeRF Model
**Redirected from the JAX repo  (https://github.com/google/jax/issues/9139#issue-1096888310)**

Tensorflow vers: 2.7; JAX vers: 0.2.24; jaxlib vers: 0.1.72+cuda111; FLAX vers: 0.3.6

The following code is based on the MNIST FLAX jax2tf example, which I adapted for JAX NeRF:

```python
import collections
from os import path
from absl import app
from absl import flags
from flax.training import checkpoints
from jax import random
from jax.experimental.jax2tf.examples import saved_model_lib
from nerf import models
from nerf import utils
import tensorflow as tf

FLAGS = flags.FLAGS

utils.define_flags()

def main(unused_argv):
    rng = random.PRNGKey(20200823)
    rng, key = random.split(rng)
    utils.update_flags(FLAGS)
    utils.check_flags(FLAGS)    
    model, state = models.get_model_state(key, FLAGS, restore=False)
    print('Loading model')
    state = checkpoints.restore_checkpoint(FLAGS.train_dir, state)
    params = state.optimizer.target
    predict_fn = lambda params, input: model.apply({"params": params}, input)
    Rays = collections.namedtuple("Rays", ("origins", "directions", "viewdirs"))
    input_signatures = [Rays(origins=tf.TensorSpec((3,),tf.float32),directions=tf.TensorSpec((3,),tf.float32),viewdirs=tf.TensorSpec((3,),tf.float32))]
    saved_model_lib.convert_and_save_model(
        predict_fn,
        params,
        '/any/path/',
        input_signatures=input_signatures) 

if __name__ == "__main__":
    app.run(main)
```
In order to simplify the inputs to the network, and since I am only interested in running inference in TF, I initialize the RNG keys and `randomized` NeRF model inputs to `None` and `False` respectively, so that only the `rays` are inputted. This is the only change over the original JAX NeRF code:

```python
def __call__(self, rays, rng_0 = None, rng_1=None, randomized=False, depth_gt = None, rgb_only = False,depth_sampling = False):
        """Nerf Model.

        Args:
          rng_0: jnp.ndarray, random number generator for coarse model sampling.
          rng_1: jnp.ndarray, random number generator for fine model sampling.
          rays: util.Rays, a namedtuple of ray origins, directions, and viewdirs.
          randomized: bool, use randomized stratified sampling.
          rgb_only: bool, return only rgb

        Returns:
          ret: list, [(rgb_coarse, disp_coarse, acc_coarse), (rgb, disp, acc)]
        """
        # Stratified sampling along rays
        if (randomized):  
            key, rng_0 = random.split(rng_0)
        else:
            key = None
```
(also, every call to `model.apply()`  has its args order inverted to match this)

The error is prompted when attempting to compute the TF graph in this line of 'saved_model_lib.py':

```python
tf_graph = tf.function(lambda inputs: tf_fn(param_vars, inputs),
                         autograph=False,
                         experimental_compile=compile_model)
```

Full error stack: 

```
Traceback (most recent call last):
  File "/home/jorge/anaconda3/envs/jaxnerf/lib/python3.8/runpy.py", line 194, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/home/jorge/anaconda3/envs/jaxnerf/lib/python3.8/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/home/jorge/jaxnerf/nerf/save_jax_as_tf.py", line 45, in <module>
    app.run(main)
  File "/home/jorge/anaconda3/envs/jaxnerf/lib/python3.8/site-packages/absl/app.py", line 312, in run
    _run_main(main, args)
  File "/home/jorge/anaconda3/envs/jaxnerf/lib/python3.8/site-packages/absl/app.py", line 258, in _run_main
    sys.exit(main(argv))
  File "/home/jorge/jaxnerf/nerf/save_jax_as_tf.py", line 38, in main
    saved_model_lib.convert_and_save_model(
  File "/home/jorge/anaconda3/envs/jaxnerf/lib/python3.8/site-packages/jax/experimental/jax2tf/examples/saved_model_lib.py", line 114, in convert_and_save_model
    tf_graph.get_concrete_function(input_signatures[0])
  File "/home/jorge/anaconda3/envs/jaxnerf/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py", line 1259, in get_concrete_function
    concrete = self._get_concrete_function_garbage_collected(*args, **kwargs)
  File "/home/jorge/anaconda3/envs/jaxnerf/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py", line 1239, in _get_concrete_function_garbage_collected
    self._initialize(args, kwargs, add_initializers_to=initializers)
  File "/home/jorge/anaconda3/envs/jaxnerf/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py", line 780, in _initialize
    self._stateful_fn._get_concrete_function_internal_garbage_collected(  # pylint: disable=protected-access
  File "/home/jorge/anaconda3/envs/jaxnerf/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 3157, in _get_concrete_function_internal_garbage_collected
    graph_function, _ = self._maybe_define_function(args, kwargs)
  File "/home/jorge/anaconda3/envs/jaxnerf/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 3557, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "/home/jorge/anaconda3/envs/jaxnerf/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 3392, in _create_graph_function
    func_graph_module.func_graph_from_py_func(
  File "/home/jorge/anaconda3/envs/jaxnerf/lib/python3.8/site-packages/tensorflow/python/framework/func_graph.py", line 1143, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/home/jorge/anaconda3/envs/jaxnerf/lib/python3.8/site-packages/tensorflow/python/eager/def_function.py", line 672, in wrapped_fn
    out = weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "/home/jorge/anaconda3/envs/jaxnerf/lib/python3.8/site-packages/jax/experimental/jax2tf/examples/saved_model_lib.py", line 107, in <lambda>
    tf_graph = tf.function(lambda inputs: tf_fn(param_vars, inputs),
  File "/home/jorge/anaconda3/envs/jaxnerf/lib/python3.8/site-packages/jax/experimental/jax2tf/jax2tf.py", line 418, in converted_fun
    out_with_avals = _interpret_fun(flat_fun, args_flat, args_avals_flat,
  File "/home/jorge/anaconda3/envs/jaxnerf/lib/python3.8/site-packages/jax/experimental/jax2tf/jax2tf.py", line 486, in _interpret_fun
    fun.call_wrapped(*in_vals)
  File "/home/jorge/anaconda3/envs/jaxnerf/lib/python3.8/site-packages/jax/linear_util.py", line 166, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/home/jorge/anaconda3/envs/jaxnerf/lib/python3.8/site-packages/jax/experimental/jax2tf/jax2tf.py", line 272, in fun_no_kwargs
    return fun(*args, **kwargs)
  File "/home/jorge/jaxnerf/nerf/save_jax_as_tf.py", line 35, in <lambda>
    predict_fn = lambda params, input: model.apply({"params": params}, input)
  File "/home/jorge/jaxnerf/nerf/nerf/models.py", line 268, in __call__
    raw_rgb, raw_sigma = self.MLP_0(samples_enc)
  File "/home/jorge/jaxnerf/nerf_sh/nerf/model_utils.py", line 70, in __call__
    x = dense_layer(self.net_width)(x)
  File "/home/jorge/anaconda3/envs/jaxnerf/lib/python3.8/site-packages/flax/linen/linear.py", line 171, in __call__
    kernel = self.param('kernel',
flax.errors.ScopeParamNotFoundError: No parameter named "kernel" exists in "/MLP_0/Dense_0". (https://flax.readthedocs.io/en/latest/flax.errors.html#flax.errors.ScopeParamNotFoundError)
```

Has anyone else attempted to save a JAX NeRF model using jax2tf and encountered any such issue?






Discussion:
This looks like the error that happens when the top-level "packaging" of the `params` pytree is off slightly - e.g. passing in the bare `params` tree or `{'params': {'params': params}` rather than `{'params': params}`...  (I say this as we're getting an error here at what looks like the very first parameter lookup into the pytree.)

Can you add these imports
```python
import jax
from jax import numpy as jnp
```
and before the final call to `saved_model_lib.convert_and_save_model` can you add a line: 
```python
print(jax.tree_map(jnp.shape, params)
```
so we can check the pytree structure of what you're passing in?

I tried setting up a repro quickly, but I'm seeing `nerf.models.get_model_state` which is a function that doesn't exist in JAX Nerf at https://github.com/google-research/google-research/blob/master/jaxnerf/nerf/models.py  -- what code are you actually using here?  More info needed for me to look into this.

Hi, @levskaya, thanks for your prompt response! I did as you said and this is what I am getting: 

```
FrozenDict({
    params: {
        MLP_0: {
            Dense_0: {
                bias: (256,),
                kernel: (63, 256),
            },
            Dense_1: {
                bias: (256,),
                kernel: (256, 256),
            },
            Dense_2: {
                bias: (256,),
                kernel: (256, 256),
            },
            Dense_3: {
                bias: (256,),
                kernel: (256, 256),
            },
            Dense_4: {
                bias: (256,),
                kernel: (256, 256),
            },
            Dense_5: {
                bias: (256,),
                kernel: (319, 256),
            },
            Dense_6: {
                bias: (256,),
                kernel: (256, 256),
            },
            Dense_7: {
                bias: (256,),
                kernel: (256, 256),
            },
            Dense_8: {
                bias: (1,),
                kernel: (256, 1),
            },
            Dense_9: {
                bias: (3,),
                kernel: (256, 3),
            },
        },
        MLP_1: {
            Dense_0: {
                bias: (256,),
                kernel: (63, 256),
            },
            Dense_1: {
                bias: (256,),
                kernel: (256, 256),
            },
            Dense_2: {
                bias: (256,),
                kernel: (256, 256),
            },
            Dense_3: {
                bias: (256,),
                kernel: (256, 256),
            },
            Dense_4: {
                bias: (256,),
                kernel: (256, 256),
            },
            Dense_5: {
                bias: (256,),
                kernel: (319, 256),
            },
            Dense_6: {
                bias: (256,),
                kernel: (256, 256),
            },
            Dense_7: {
                bias: (256,),
                kernel: (256, 256),
            },
            Dense_8: {
                bias: (1,),
                kernel: (256, 1),
            },
            Dense_9: {
                bias: (3,),
                kernel: (256, 3),
            },
        },
    },
})

```

Also, as you pointed out `nerf.models.get_model_state` is not an existing function in JAX NeRF; it's just a small helper, sorry I didn't include it in the first place. It's this function:

```python
def get_model_state(key, args, restore=True):
    """
    Helper for loading model with get_model & creating optimizer &
    optionally restoring checkpoint to reduce boilerplate
    """
    model, variables = get_model(key, args)
    optimizer = flax.optim.Adam(args.lr_init).create(variables)
    state = utils.TrainState(optimizer=optimizer)
    if restore:
        from flax.training import checkpoints
        state = checkpoints.restore_checkpoint(args.train_dir, state)
    return model, state
```

You should be able to reproduce with this. I am using a slightly different code to JAX NeRF but I was able to reproduce with this and their code.

Just in case it helps, I will add some more info about my setup:

Ubuntu 20.04 on WSL2
RTX 3080 (CUDA 11.2 CUDDN 8.1.1 NVIDIA driver 510.06)
Tensorflow 2.7
Jaxlib 0.1.74+cuda11.cudnn805 (upgraded a few hours ago but same result)
Jax 0.2.26
Flax 0.3.6

**EDIT:**  I was also able to reproduce it in the following setup:
Native Ubuntu 18.04
RTX 2080ti (CUDA 10.1 CUDDN 7.6.5 NVIDIA driver 418.87.00)
Tensorflow 2.3.1
Jaxlib 0.1.72+cuda111 
Jax 0.2.26
Flax 0.3.6


And what about `utils.check_flags(FLAGS)` (also doesn't exist in original repo)  and `utils.update_flags(FLAGS)`?  I have no idea what config you're actually running here?  Is there a link to your code and whatever FLAGS are actually being used?
If `get_model(key, args)` is `jaxnerf.nerf.models.get_model` it doesn't have the right signature.  

Did you define this helper as well?  How are you specifying the example_batch that `jaxnerf.nerf.models.get_model` needs, which is ultimately calling `jaxnerf.nerf.models.construct_nerf(key, example_batch, args) --> model, init_variables`
Sorry, one more question - what do you mean you "initialize the RNG keys" to `None`??  you can't just set things like `key` `rng_0` and `rng_1` to `None` in the original JAX Nerf code... those are JAX deterministic PRNG keys that have to be provided.  

If you have something running at all you must have heavily altered the original `__call__` function - I really need to see your code to have any idea about what's going on here.  Please just dump all your changes somewhere so I can see what's actually being run.  

A quick guess is that you accidentally changed the nested module structure which is causing a mismatch between the provided parameters and the model structure.
Hi, a few things regarding your comments:

1) For reproducibility, I used the lego config from the jaxnerf code. This is what `utlis.update_flags(FLAGS)`  loads. This gets me the following params map:

```
FrozenDict({
    params: {
        MLP_0: {
            Dense_0: {
                bias: (256,),
                kernel: (63, 256),
            },
            Dense_1: {
                bias: (256,),
                kernel: (256, 256),
            },
            Dense_10: {
                bias: (128,),
                kernel: (283, 128),
            },
            Dense_11: {
                bias: (3,),
                kernel: (128, 3),
            },
            Dense_2: {
                bias: (256,),
                kernel: (256, 256),
            },
            Dense_3: {
                bias: (256,),
                kernel: (256, 256),
            },
            Dense_4: {
                bias: (256,),
                kernel: (256, 256),
            },
            Dense_5: {
                bias: (256,),
                kernel: (319, 256),
            },
            Dense_6: {
                bias: (256,),
                kernel: (256, 256),
            },
            Dense_7: {
                bias: (256,),
                kernel: (256, 256),
            },
            Dense_8: {
                bias: (1,),
                kernel: (256, 1),
            },
            Dense_9: {
                bias: (256,),
                kernel: (256, 256),
            },
        },
        MLP_1: {
            Dense_0: {
                bias: (256,),
                kernel: (63, 256),
            },
            Dense_1: {
                bias: (256,),
                kernel: (256, 256),
            },
            Dense_10: {
                bias: (128,),
                kernel: (283, 128),
            },
            Dense_11: {
                bias: (3,),
                kernel: (128, 3),
            },
            Dense_2: {
                bias: (256,),
                kernel: (256, 256),
            },
            Dense_3: {
                bias: (256,),
                kernel: (256, 256),
            },
            Dense_4: {
                bias: (256,),
                kernel: (256, 256),
            },
            Dense_5: {
                bias: (256,),
                kernel: (319, 256),
            },
            Dense_6: {
                bias: (256,),
                kernel: (256, 256),
            },
            Dense_7: {
                bias: (256,),
                kernel: (256, 256),
            },
            Dense_8: {
                bias: (1,),
                kernel: (256, 1),
            },
            Dense_9: {
                bias: (256,),
                kernel: (256, 256),
            },
        },
    },
})
```

2) `utils.check_flags(FLAGS)` does indeed not exist in the original jaxnerf, sorry about that. It's just a helper to check whether the user has set training and data dirs. Can be removed without issue; this code also reproduces the error:

```python
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"

import collections
from os import path

from absl import app
from absl import flags
import jax
from jax import numpy as jnp
from jax import random
from jax.experimental.jax2tf.examples import saved_model_lib

from nerf_sh.nerf import models
from nerf_sh.nerf import utils
import tensorflow as tf

FLAGS = flags.FLAGS

utils.define_flags()

def main(unused_argv):
    rng = random.PRNGKey(20200823)
    rng, key = random.split(rng)

    utils.update_flags(FLAGS) 
    model, state = models.get_model_state(key, FLAGS, restore=True)
    params = state.optimizer.target
    predict_fn = lambda params, input: model.apply({"params": params}, input)
    Rays = collections.namedtuple("Rays", ("origins", "directions", "viewdirs"))
    input_signatures = [Rays(origins=tf.TensorSpec((3,),tf.float32),directions=tf.TensorSpec((3,),tf.float32),viewdirs=tf.TensorSpec((3,),tf.float32))]
    print(jax.tree_map(jnp.shape, params))
    saved_model_lib.convert_and_save_model(
        predict_fn,
        params,
        '/some/path',
        input_signatures=input_signatures) 

if __name__ == "__main__":
    app.run(main)
```

3) `get_model(key,args)` is slightly changed from JAX NeRF, as I didn't need the dataset peeking func.  However if you call the function as `model, variables = models.get_model(key, dataset.peek(), FLAGS)` after loading the dataset with `dataset = datasets.get_dataset("test", FLAGS)` you should be able to reproduce as well.  In any case, this is the function I am using:

```python
def get_model(key, args):
    """A helper function that wraps around a 'model zoo'."""
    model_dict = {
        "nerf": construct_nerf,
    }
    return model_dict[args.model](key, args)
```

4) The JAX deterministic PRNG keys can be initialized to None as long as you use the model for test inference: these keys are only used in the randomized stratified sampling scheme which is only performed during training; for test rendering, sampling becomes deterministic. Thus, we can alter the order of the `__call__` parameters and input a single tuple of raydirs, origins and viewdirs. The new call function looks like this: 

```python
def __call__(self, rays, rng_0 = None, rng_1=None, randomized=False):
        """Nerf Model.

        Args:
          rng_0: jnp.ndarray, random number generator for coarse model sampling.
          rng_1: jnp.ndarray, random number generator for fine model sampling.
          rays: util.Rays, a namedtuple of ray origins, directions, and viewdirs.
          randomized: bool, use randomized stratified sampling.

        Returns:
          ret: list, [(rgb_coarse, disp_coarse, acc_coarse), (rgb, disp, acc)]
        """
        # Stratified sampling along rays
        if (randomized):  
            key, rng_0 = random.split(rng_0)
        else:
            key = None
        z_vals, samples = model_utils.sample_along_rays(
            key,
            rays.origins,
            rays.directions,
            self.num_coarse_samples,
            self.near,
            self.far,
            randomized,
            self.lindisp
        )

        samples_enc = model_utils.posenc(
            samples,
            self.min_deg_point,
            self.max_deg_point,
            self.legacy_posenc_order,
        )

        # Point attribute predictions
        if self.use_viewdirs:
            viewdirs_enc = model_utils.posenc(
                rays.viewdirs,
                0,
                self.deg_view,
                self.legacy_posenc_order,
            )
            raw_rgb, raw_sigma = self.MLP_0(samples_enc, viewdirs_enc)
        else:
            raw_rgb, raw_sigma = self.MLP_0(samples_enc)
        # Add noises to regularize the density predictions if needed
        key, rng_0 = random.split(rng_0)
        raw_sigma = model_utils.add_gaussian_noise(
            key,
            raw_sigma,
            self.noise_std,
            randomized,
        )

        rgb = self.rgb_activation(raw_rgb)
        sigma = self.sigma_activation(raw_sigma)

        comp_rgb, disp, acc, weights,depth = model_utils.volumetric_rendering(
            rgb,
            sigma,
            z_vals,
            rays.directions,
            white_bkgd=self.white_bkgd,
        )
        ret = [
            (comp_rgb, disp, acc,depth),
        ]
        # Hierarchical sampling based on coarse predictions
        if self.num_fine_samples > 0:
            z_vals_mid = 0.5 * (z_vals[Ellipsis, 1:] + z_vals[Ellipsis, :-1])
            if (randomized): key, rng_1 = random.split(rng_1)
            z_vals, samples = model_utils.sample_pdf(
                key,
                z_vals_mid,
                weights[Ellipsis, 1:-1],
                rays.origins,
                rays.directions,
                z_vals,
                self.num_fine_samples,
                randomized,
            )
            samples_enc = model_utils.posenc(
                samples,
                self.min_deg_point,
                self.max_deg_point,
                self.legacy_posenc_order,
            )

            if self.use_viewdirs:
                raw_rgb, raw_sigma = self.MLP_1(samples_enc, viewdirs_enc)
            else:
                raw_rgb, raw_sigma = self.MLP_1(samples_enc)
            if (randomized): key, rng_1 = random.split(rng_1)
            raw_sigma = model_utils.add_gaussian_noise(
                key,
                raw_sigma,
                self.noise_std,
                randomized,
            )
         
            rgb = self.rgb_activation(raw_rgb)
            sigma = self.sigma_activation(raw_sigma)

            comp_rgb, disp, acc, unused_weights, depth = model_utils.volumetric_rendering(
                rgb,
                sigma,
                z_vals,
                rays.directions,
                white_bkgd=self.white_bkgd,
            )
            ret.append((comp_rgb, disp, acc,depth))
        return ret
```

Then the only change you need to make is in the args order of the `model.apply()` calls. In `train.py`, function `loss_fn()` you replace `ret = model.apply(variables, key_0, key_1, rays, FLAGS.randomized)` with `ret = model.apply(variables, rays, key_0, key_1, FLAGS.randomized)`. You do the same in `train.py`, `render_fn()`

That should be it. Lemme know if I missed something!




@Arcanous98 - thanks for providing the extra info!

Actually I just noticed something from your first response that I should have noticed immediately:

if the output of the inserted printout:
```python
input_signatures = [Rays(origin....
print(jax.tree_map(jnp.shape, params)
saved_model_lib.convert_and_save_model(...
```

has this structure:
```python
FrozenDict({
    params: {
        MLP_0: { ... }
        ...
    }
})
```

The `params` object shouldn't itself have an extra `params:` layer inside it, since in your `predict_fn` function you write:
```python
predict_fn = lambda params, input: model.apply({"params": params}, input)
```

which is adding an extra nesting layer under another `"params"` key, which would lead to precisely the error that you're seeing.  

In Flax the `init` function returns, and the `apply` function takes a variable (frozen-) dictionary structured at the top-level like 

```python
{
 "params": nested_param_dict, 
 "some_stateful_collection": some_stateful_collection_dict, 
 "some_other_stateful_collection": some_other_stateful_collection_dict
...
}
```
where each of those nested_dicts share the same module-defined nesting structure.

If you try to remove the extra `{"params": ...}` nesting, does your code run correctly?

It does! Thanks for the help @levskaya. I'm closing the issue now 👍 
Great!  Happy to Help!  Happy NeRFing. ;)

Resolving patch:
diff --git a/flax/core/scope.py b/flax/core/scope.py
--- a/flax/core/scope.py
+++ b/flax/core/scope.py
@@ -769,7 +769,7 @@ def bind(variables: VariableDict,
   across the JAX software ecosystem.
   """
   if not _is_valid_variables(variables):
-    raise errors.ApplyScopeInvalidVariablesError()
+    raise errors.ApplyScopeInvalidVariablesTypeError()
   if rngs is not None and not _is_valid_rngs(rngs):
     raise errors.InvalidRngError(
       'rngs should be a dictionary mapping strings to `jax.PRNGKey`.')
@@ -794,6 +794,12 @@ def wrapper(variables: VariableDict,
               *args,
               rngs: Optional[RNGSequences] = None,
               **kwargs) -> Union[Any, Tuple[Any, VariableDict]]:
+    # Try to detect if user accidentally passed {'params': {'params': ...}.
+    if 'params' in variables and isinstance(
+        variables['params'], 
+        (dict, FrozenDict)) and 'params' in variables['params']:
+      raise errors.ApplyScopeInvalidVariablesStructureError(variables)
+
     with bind(variables, rngs=rngs, mutable=mutable).temporary() as root:
       y = fn(root, *args, **kwargs)
     if mutable is not False:
diff --git a/flax/errors.py b/flax/errors.py
--- a/flax/errors.py
+++ b/flax/errors.py
@@ -122,7 +122,7 @@ def __init__(self, msg):
     super().__init__(msg)
 
 
-class ApplyScopeInvalidVariablesError(FlaxError):
+class ApplyScopeInvalidVariablesTypeError(FlaxError):
   """
   When calling :meth:`Module.apply() <flax.linen.Module.apply>`, the first
   argument should be a variable dict. For more explanation on variable dicts,
@@ -134,6 +134,18 @@ def __init__(self):
                      'dictionary with string keys.')
 
 
+class ApplyScopeInvalidVariablesStructureError(FlaxError):
+  """
+  This error is thrown when the dict passed as `variables` to apply() has an
+  extra 'params' layer, i.e. {'params': {'params': ...}}.
+  For more explanation on variable dicts, please see :mod:`flax.core.variables`.
+  """
+  def __init__(self, variables):
+    super().__init__(f'Expected the first argument passed to an apply function '
+                     'to be a dictionary containing a \'params\' key at the '
+                     'root level, but got "{variables}".')
+
+
 class ScopeParamNotFoundError(FlaxError):
   """
   This error is thrown when trying to access a parameter that does not exist.
@@ -176,7 +188,7 @@ class ScopeCollectionNotFound(FlaxError):
   def __init__(self, col_name, var_name, scope_path):
     super().__init__(
       f'Tried to access "{var_name}" from collection "{col_name}"" in '
-      f'"{scope_path}" but the collection is emtpy.')
+      f'"{scope_path}" but the collection is empty.')
 
 
 class ScopeParamShapeError(FlaxError):


Test patch:
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
@@ -111,6 +111,21 @@ def f(scope):
     with self.assertRaisesRegex(errors.ScopeParamShapeError, msg):
       apply(f)(freeze({'params': {'test': np.ones((2,))}}))
 
+  def test_apply_variables_bad_pytree(self):
+    def f(scope):
+      scope.param('kernel', nn.initializers.ones, (4,))
+
+    params = freeze({
+        'params': {
+            'kernel': np.ones((4,)),
+        },
+    })
+    apply(f)(params)  # Valid.
+    msg = 'dictionary containing a \'params\' key at the root level'
+    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)