Repository: google/flax

Issue:
Linen: cryptic error message when feeding with incorrect rngs keys
I mention a problem I encountered recently, it cost me a lot of time since the error message is cryptic

### Problem you have encountered:

I made a typing mistake (shame on me) : 
``` python 

# a good key
key1, key2,key3  = random.split(random.PRNGKey(0), 3)

# mistake while typing
bad_key = random.split(random.PRNGKey(0), 2)

```
And then cryptic message in `init` or `apply`: 

```python
m = MyModule()
p = m.init({'params':key1,'dropout':bad_key},x)
```


### Steps to reproduce:

https://colab.research.google.com/drive/1Ijr74leHGN8ZrvipgpQnVo9Ql8SI03-Y?usp=sharing

### Logs, error messages, etc:

```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-6-e5d297b6aa21> in <module>()
----> 1 p = m.init({'params':key1,'dropout':bad_key},x)

34 frames
/usr/local/lib/python3.6/dist-packages/flax/linen/module.py in init(self, rngs, method, *args, **kwargs)
    474   def init(self, rngs, *args, method=None, **kwargs):
    475     """Create and return initialized data for module with rngs."""
--> 476     _, v_out = self.init_with_output(rngs, *args, method=method, **kwargs)
    477     return v_out
    478 

/usr/local/lib/python3.6/dist-packages/flax/linen/module.py in init_with_output(self, rngs, method, *args, **kwargs)
    470       rngs = {'params': rngs}
    471     return self.apply(
--> 472         {}, *args, rngs=rngs, method=method, mutable=True, **kwargs)
    473 
    474   def init(self, rngs, *args, method=None, **kwargs):

/usr/local/lib/python3.6/dist-packages/flax/linen/module.py in apply(self, variables, rngs, method, mutable, *args, **kwargs)
    462     fn = lambda scope: method(self.clone(parent=scope),
    463                               *args, **kwargs)
--> 464     return apply(fn, mutable=mutable)(variables, rngs=rngs)
    465 
    466   def init_with_output(self, rngs, *args, method=None, **kwargs):

/usr/local/lib/python3.6/dist-packages/flax/core/scope.py in wrapper(variables, rngs, *args, **kwargs)
    338     new_variables = _unfreeze_variables(variables, mutable)
    339     with Scope(new_variables, rngs=rngs).temporary() as root:
--> 340       y = fn(root, *args, **kwargs)
    341     if mutable:
    342       return y, freeze(new_variables)

/usr/local/lib/python3.6/dist-packages/flax/linen/module.py in <lambda>(scope)
    461       method = get_unbound_fn(method)
    462     fn = lambda scope: method(self.clone(parent=scope),
--> 463                               *args, **kwargs)
    464     return apply(fn, mutable=mutable)(variables, rngs=rngs)
    465 

/usr/local/lib/python3.6/dist-packages/flax/linen/module.py in wrapped_module_method(self, *args, **kwargs)
    154     _context.module_stack.append(self)
    155     try:
--> 156       return fun(self, *args, **kwargs)
    157     finally:
    158       _context.module_stack.pop()

<ipython-input-3-efadaf5263bf> in __call__(self, x)
      3     @nn.compact
      4     def __call__(self, x):
----> 5         self.make_rng('dropout')
      6         return x

/usr/local/lib/python3.6/dist-packages/flax/linen/module.py in make_rng(self, kind)
    451   def make_rng(self, kind: str) -> PRNGKey:
    452     """Get a new rng key of a given kind from this Module."""
--> 453     return self.scope.make_rng(kind)
    454 
    455   def apply(self, variables, *args, rngs=None,

/usr/local/lib/python3.6/dist-packages/flax/core/scope.py in make_rng(self, name)
    272     self._validate_trace_level()
    273     self.rng_counters[name] += 1
--> 274     return random.fold_in(self.rngs[name], self.rng_counters[name])
    275 
    276   def get_variable(self, col: str, name: str, default: T = None) -> T:

/usr/local/lib/python3.6/dist-packages/jax/random.py in fold_in(key, data)
    294     statistically safe for producing a stream of new pseudo-random values.
    295   """
--> 296   return _fold_in(key, data)
    297 
    298 @jit

/usr/local/lib/python3.6/dist-packages/jax/api.py in f_jitted(*args, **kwargs)
    213         backend=backend,
    214         name=flat_fun.__name__,
--> 215         donated_invars=donated_invars)
    216     return tree_unflatten(out_tree(), out)
    217 

/usr/local/lib/python3.6/dist-packages/jax/core.py in bind(self, fun, *args, **params)
   1142 
   1143   def bind(self, fun, *args, **params):
-> 1144     return call_bind(self, fun, *args, **params)
   1145 
   1146   def process(self, trace, fun, tracers, params):

/usr/local/lib/python3.6/dist-packages/jax/core.py in call_bind(primitive, fun, *args, **params)
   1133   tracers = map(top_trace.full_raise, args)
   1134   with maybe_new_sublevel(top_trace):
-> 1135     outs = primitive.process(top_trace, fun, tracers, params)
   1136   return map(full_lower, apply_todos(env_trace_todo(), outs))
   1137 

/usr/local/lib/python3.6/dist-packages/jax/core.py in process(self, trace, fun, tracers, params)
   1145 
   1146   def process(self, trace, fun, tracers, params):
-> 1147     return trace.process_call(self, fun, tracers, params)
   1148 
   1149   def post_process(self, trace, out_tracers, params):

/usr/local/lib/python3.6/dist-packages/jax/core.py in process_call(self, primitive, f, tracers, params)
    575 
    576   def process_call(self, primitive, f, tracers, params):
--> 577     return primitive.impl(f, *tracers, **params)
    578   process_map = process_call
    579 

/usr/local/lib/python3.6/dist-packages/jax/interpreters/xla.py in _xla_call_impl(fun, device, backend, name, donated_invars, *args)
    528 def _xla_call_impl(fun: lu.WrappedFun, *args, device, backend, name, donated_invars):
    529   compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,
--> 530                                *unsafe_map(arg_spec, args))
    531   try:
    532     return compiled_fun(*args)

/usr/local/lib/python3.6/dist-packages/jax/linear_util.py in memoized_fun(fun, *args)
    232       fun.populate_stores(stores)
    233     else:
--> 234       ans = call(fun, *args)
    235       cache[key] = (ans, fun.stores)
    236     return ans

/usr/local/lib/python3.6/dist-packages/jax/interpreters/xla.py in _xla_callable(fun, device, backend, name, donated_invars, *arg_specs)
    593   abstract_args, arg_devices = unzip2(arg_specs)
    594   if config.omnistaging_enabled:
--> 595     jaxpr, out_avals, consts = pe.trace_to_jaxpr_final(fun, abstract_args)
    596     if any(isinstance(c, core.Tracer) for c in consts):
    597       raise core.UnexpectedTracerError("Encountered an unexpected tracer.")

/usr/local/lib/python3.6/dist-packages/jax/interpreters/partial_eval.py in trace_to_jaxpr_final(fun, in_avals)
   1021     main.source_info = fun_sourceinfo(fun.f)  # type: ignore
   1022     main.jaxpr_stack = ()  # type: ignore
-> 1023     jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, main, in_avals)
   1024     del main
   1025   return jaxpr, out_avals, consts

/usr/local/lib/python3.6/dist-packages/jax/interpreters/partial_eval.py in trace_to_subjaxpr_dynamic(fun, main, in_avals)
   1002     trace = DynamicJaxprTrace(main, core.cur_sublevel())
   1003     in_tracers = map(trace.new_arg, in_avals)
-> 1004     ans = fun.call_wrapped(*in_tracers)
   1005     out_tracers = map(trace.full_raise, ans)
   1006   jaxpr, out_avals, consts = frame.to_jaxpr(in_tracers, out_tracers)

/usr/local/lib/python3.6/dist-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    149 
    150     try:
--> 151       ans = self.f(*args, **dict(self.params, **kwargs))
    152     except:
    153       # Some transformations yield from inside context managers, so we have to

/usr/local/lib/python3.6/dist-packages/jax/random.py in _fold_in(key, data)
    298 @jit
    299 def _fold_in(key, data):
--> 300   return threefry_2x32(key, PRNGKey(data))
    301 
    302 

/usr/local/lib/python3.6/dist-packages/jax/api.py in f_jitted(*args, **kwargs)
    213         backend=backend,
    214         name=flat_fun.__name__,
--> 215         donated_invars=donated_invars)
    216     return tree_unflatten(out_tree(), out)
    217 

/usr/local/lib/python3.6/dist-packages/jax/core.py in bind(self, fun, *args, **params)
   1142 
   1143   def bind(self, fun, *args, **params):
-> 1144     return call_bind(self, fun, *args, **params)
   1145 
   1146   def process(self, trace, fun, tracers, params):

/usr/local/lib/python3.6/dist-packages/jax/core.py in call_bind(primitive, fun, *args, **params)
   1133   tracers = map(top_trace.full_raise, args)
   1134   with maybe_new_sublevel(top_trace):
-> 1135     outs = primitive.process(top_trace, fun, tracers, params)
   1136   return map(full_lower, apply_todos(env_trace_todo(), outs))
   1137 

/usr/local/lib/python3.6/dist-packages/jax/core.py in process(self, trace, fun, tracers, params)
   1145 
   1146   def process(self, trace, fun, tracers, params):
-> 1147     return trace.process_call(self, fun, tracers, params)
   1148 
   1149   def post_process(self, trace, out_tracers, params):

/usr/local/lib/python3.6/dist-packages/jax/interpreters/partial_eval.py in process_call(self, call_primitive, f, tracers, params)
    938   def process_call(self, call_primitive, f, tracers, params):
    939     in_avals = [t.aval for t in tracers]
--> 940     jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(f, self.main, in_avals)
    941     if not jaxpr.eqns:
    942       return core.eval_jaxpr(jaxpr, consts, *tracers)

/usr/local/lib/python3.6/dist-packages/jax/interpreters/partial_eval.py in trace_to_subjaxpr_dynamic(fun, main, in_avals)
   1002     trace = DynamicJaxprTrace(main, core.cur_sublevel())
   1003     in_tracers = map(trace.new_arg, in_avals)
-> 1004     ans = fun.call_wrapped(*in_tracers)
   1005     out_tracers = map(trace.full_raise, ans)
   1006   jaxpr, out_avals, consts = frame.to_jaxpr(in_tracers, out_tracers)

/usr/local/lib/python3.6/dist-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    149 
    150     try:
--> 151       ans = self.f(*args, **dict(self.params, **kwargs))
    152     except:
    153       # Some transformations yield from inside context managers, so we have to

/usr/local/lib/python3.6/dist-packages/jax/random.py in threefry_2x32(keypair, count)
    261   out = jnp.concatenate(x)
    262   assert out.dtype == np.uint32
--> 263   return lax.reshape(out[:-1] if odd_size else out, count.shape)
    264 
    265 

/usr/local/lib/python3.6/dist-packages/jax/lax/lax.py in reshape(operand, new_sizes, dimensions)
    688     return reshape_p.bind(
    689       operand, new_sizes=new_sizes,
--> 690       dimensions=None if dimensions is None or same_dims else tuple(dimensions))
    691 
    692 def pad(operand: Array, padding_value: Array,

/usr/local/lib/python3.6/dist-packages/jax/core.py in bind(self, *args, **params)
    264     top_trace = find_top_trace(args)
    265     tracers = map(top_trace.full_raise, args)
--> 266     out = top_trace.process_primitive(self, tracers, params)
    267     return map(full_lower, out) if self.multiple_results else full_lower(out)
    268 

/usr/local/lib/python3.6/dist-packages/jax/interpreters/partial_eval.py in process_primitive(self, primitive, tracers, params)
    926   def process_primitive(self, primitive, tracers, params):
    927     avals = [t.aval for t in tracers]
--> 928     out_avals = primitive.abstract_eval(*avals, **params)
    929     out_avals = [out_avals] if not primitive.multiple_results else out_avals
    930     source_info = source_info_util.current()

/usr/local/lib/python3.6/dist-packages/jax/lax/lax.py in standard_abstract_eval(prim, shape_rule, dtype_rule, *args, **kwargs)
   1909     return ConcreteArray(prim.impl(*[x.val for x in args], **kwargs))
   1910   elif least_specialized is ShapedArray:
-> 1911     return ShapedArray(shape_rule(*args, **kwargs), dtype_rule(*args, **kwargs))
   1912   elif least_specialized is UnshapedArray:
   1913     return UnshapedArray(dtype_rule(*args, **kwargs))

/usr/local/lib/python3.6/dist-packages/jax/lax/lax.py in _reshape_shape_rule(operand, new_sizes, dimensions)
   3365   if prod(np.shape(operand)) != prod(new_sizes):
   3366     msg = 'reshape total size must be unchanged, got new_sizes {} for shape {}.'
-> 3367     raise TypeError(msg.format(new_sizes, np.shape(operand)))
   3368   if dimensions is not None:
   3369     if set(dimensions) != set(range(np.ndim(operand))):

TypeError: reshape total size must be unchanged, got new_sizes (2,) for shape (4,).
```

### Steps to reproduce:

https://colab.research.google.com/drive/1Ijr74leHGN8ZrvipgpQnVo9Ql8SI03-Y?usp=sharing


Discussion:
Sure, it would be good to add assertions that arguments that we expect to be RNGs are indeed RNGs. (I see that `jax.random` has `_is_prng_key` but that's a private method, so we should either ask the JAX core folks to make this public or, as a first step, replicate it in Flax).
I'm looking into hardening init/apply arg validation (also for Frozen vs normal dict). I'll make sure the RNGs are validated as well.

Resolving patch:
diff --git a/flax/core/lift.py b/flax/core/lift.py
--- a/flax/core/lift.py
+++ b/flax/core/lift.py
@@ -70,7 +70,8 @@ def _dup_scopes(orig_scopes, scopes, paths):
 def pack(fn: Callable[..., Any],
          in_variable_filters: Sequence[CollectionFilter],
          out_variable_filters: Sequence[CollectionFilter],
-         rng_filters: Sequence[PRNGSequenceFilter]) -> Callable[..., Any]:
+         rng_filters: Sequence[PRNGSequenceFilter],
+         name=None) -> Callable[..., Any]:
   """Pack variables and rngs for functional transformations.
 
   The pack function is the building block for all other lifted transformations.
@@ -123,9 +124,16 @@ def scope_fn(variable_groups_xs, rng_groups_xs):
         # make sure variable dicts are cloned and can't be manipulated by ref sharing.
         variables = jax.tree_map(lambda x: x, variables)
         scope_mutable = intersect_filters(scope.root.mutable, mutable)
+        new_path = scope.path
+        if name:
+          if new_path:
+            new_path = new_path[:-1] + (f'{name}({new_path[-1]})',)
+          else:
+            new_path = (f'{name}()',)
         inner_scope = Scope(
             variables, name=scope.name, rngs=rngs,
-            mutable=scope_mutable, parent=None)
+            mutable=scope_mutable, parent=None,
+            path=new_path)
         inner_scopes.append(inner_scope)
       inner_scopes = _dup_scopes(scopes, inner_scopes, paths)
       return treedef.unflatten(inner_scopes)
@@ -158,8 +166,8 @@ def repack(inner_scope_tree):
     for scope, out_variable_groups in zip(scopes, out_variable_groups_xs):
       for out_variable_group in out_variable_groups:
         for col_name, collection in out_variable_group.items():
-          for name, value in collection.items():
-            scope.put_variable(col_name, name, value)
+          for var_name, value in collection.items():
+            scope.put_variable(col_name, var_name, value)
     return y
   return wrapper
 
@@ -205,7 +213,7 @@ def wrapper(scope_fn, repack, variable_groups_xs, rng_groups_xs, fn, *args):
   is_target_out = mutable or init
   in_vars = (target, variables)
   out_vars = (target, variables) if is_target_out else ((), variables)
-  wrapper = pack(wrapper, in_vars, out_vars, (rngs,))
+  wrapper = pack(wrapper, in_vars, out_vars, (rngs,), name='transform')
   return wrapper
 
 
@@ -350,7 +358,8 @@ def mapped(variable_groups_xs, rng_groups_xs, args):
     return mapped(variable_groups_xs, rng_groups_xs, args)
 
   return pack(
-      inner, variable_in_groups, variable_out_groups, rng_groups)
+      inner, variable_in_groups, variable_out_groups, rng_groups,
+      name='vmap')
 
 
 ScanAxis = int
@@ -491,7 +500,8 @@ def scanned(broadcast_vars, carry, variable_groups_xs, rng_groups_xs, args):
       inner,
       (variable_broadcast, variable_carry) + variable_in_groups,
       (variable_broadcast, variable_carry) + variable_out_groups,
-      rng_groups)
+      rng_groups,
+      name='scan')
 
 
 def custom_vjp(fn: Callable[..., Any], backward_fn: Callable[..., Any],
@@ -560,7 +570,8 @@ def f_bwd(*args):
   variable_out_groups = (grad_kind, True,)
   rng_groups = (True,)
   return pack(
-      inner, variable_in_groups, variable_out_groups, rng_groups)
+      inner, variable_in_groups, variable_out_groups, rng_groups,
+      name='custom_vjp')
 
 
 def remat(fn: Callable[..., Any],
@@ -576,7 +587,7 @@ def rematted(variable_groups_xs, rng_groups_xs, *args):
       return y, repack_fn(scope)
 
     return rematted(variable_groups, rng_groups, *args)
-  return pack(inner, (variables,), (variables,), (rngs,))
+  return pack(inner, (variables,), (variables,), (rngs,), name='remat')
 
 
 def jit(fn: Callable[..., Any],
@@ -601,7 +612,7 @@ def jitted(variable_groups_xs, rng_groups_xs, *args):
 
     return jitted(variable_groups_xs, rng_groups_xs, *args)
 
-  return pack(inner, (variables,), (variables,), (rngs,))
+  return pack(inner, (variables,), (variables,), (rngs,), name='jit')
 
 
 def remat_scan(body_fn: Callable[..., Any], scope: Scope, carry: Any,
diff --git a/flax/core/scope.py b/flax/core/scope.py
--- a/flax/core/scope.py
+++ b/flax/core/scope.py
@@ -45,6 +45,8 @@
 PRNGKey = Any
 Array = Any
 
+RNGSequences = Dict[str, PRNGKey]
+
 Filter = Union[bool, str, Sequence[str]]
 CollectionFilter = Filter
 PRNGSequenceFilter = Filter
@@ -54,6 +56,7 @@
 MaybeFrozenCollection = Union[MutableCollection, FrozenCollection]
 
 Variables = Dict[str, MaybeFrozenCollection]
+FrozenVariables = Dict[str, FrozenCollection]
 
 
 def _fold_in_str(rng: PRNGKey, data: str) -> PRNGKey:
@@ -204,7 +207,8 @@ def __init__(self,
                rngs: Optional[Dict[str, PRNGKey]] = None,
                name: Optional[str] = None,
                mutable: CollectionFilter = False,
-               parent: Optional['Scope'] = None):
+               parent: Optional['Scope'] = None,
+               path: Tuple[str] = ()):
     """Initializes a Scope.
 
     Args:
@@ -216,6 +220,7 @@ def __init__(self,
     self._variables = variables
     self.parent = parent
     self.name = name
+    self.path = path
     self.rngs = rngs if rngs else {}
     self.mutable = mutable
 
@@ -229,6 +234,12 @@ def __init__(self,
 
     self._invalid = False
 
+
+  @property
+  def path_text(self) -> str:
+    """Returns the path as a human readable string with slashes between parts."""
+    return '/' + '/'.join(self.path)
+
   @property
   def invalid(self) -> bool:
     """Returns true if this scope is invalidated as a result of `Scope.temporary`."""
@@ -279,6 +290,8 @@ def reserve(self, name: str):
     Args:
       name: The name to reserve.
     """
+    if not isinstance(name, str):
+      raise ValueError('Variable and child scopes should have a string name.')
     if name in self.reservations:
       raise ValueError(f'Duplicate use of name: "{name}"')
     self.reservations.add(name)
@@ -315,7 +328,7 @@ def push(self, name: Optional[str] = None, prefix: str = '', reuse=False) -> 'Sc
       return self._children[name]
     self.reserve(name)
     rngs = {key: _fold_in_str(rng, name) for key, rng in self.rngs.items()}
-    scope = Scope({}, name=name, rngs=rngs, parent=self)
+    scope = Scope({}, name=name, rngs=rngs, parent=self, path=self.path + (name,))
     self._children[name] = scope
     return scope
 
@@ -358,7 +371,6 @@ def is_mutable_collection(self, col: str) -> bool:
     """Check whether a collection is mutable."""
     return in_filter(self.root.mutable, col)
 
-
   def _mutable_collection(self, col: str) -> MutableCollection:
     if not self.is_mutable_collection(col):
       raise ValueError(f'Collection is not mutable: "{col}"')
@@ -413,6 +425,10 @@ def put_variable(self, col: str, name: str, value: Any):
     """Update the value of a Variable."""
     self._check_valid()
     self._validate_trace_level()
+    if not self.is_mutable_collection(col):
+      raise ValueError(
+        f'Trying to update variable "{name}" in "{self.path_text}" '
+        f'but collection "{col}" is immutable.')
     variables = self._mutable_collection(col)
     variables[name] = value
 
@@ -421,6 +437,8 @@ def variable(self, col: str, name: str, init_fn: Callable[..., T],
     """Create a Variable."""
     self.reserve(name)
     if not self.has_variable(col, name):
+      if not self.is_mutable_collection('params'):
+        raise ValueError(f'No paramater named "{name}" exists in "{self.path_text}".')
       init_value = init_fn(*init_args)
       self.put_variable(col, name, init_value)
     return Variable(self, col, name)
@@ -441,9 +459,11 @@ def param(self, name: str, init_fn: Callable[..., T], *init_args) -> T:
         # we might intentionally change the dtype for inference to a half float type for example.
         if jnp.shape(val) != jnp.shape(abs_val):
           raise ValueError('Inconsistent shapes between value and initializer '
-                           f'for parameter "{name}": {jnp.shape(val)}, {jnp.shape(abs_val)}')
+                           f'for parameter "{name}" in "{self.path_text}": {jnp.shape(val)}, {jnp.shape(abs_val)}')
       return value
     else:
+      if not self.is_mutable_collection('params'):
+        raise ValueError(f'No paramater named "{name}" exists in "{self.path_text}".')
       value = init_fn(self.make_rng('params'), *init_args)
       self.put_variable('params', name, value)
       return value
@@ -474,7 +494,15 @@ def apply(fn: Callable[..., Any],
     `fn` with the scope partially applied.
   """
   @functools.wraps(fn)
-  def wrapper(variables, *args, rngs=None, **kwargs):
+  def wrapper(variables: FrozenVariables, *args,
+              rngs: Optional[RNGSequences] = None, **kwargs) -> (Any, FrozenVariables):
+    
+    if not _is_valid_variables(variables):
+      raise ValueError('The first argument passed to an apply function '
+                       'should be a dictionary of collections. '
+                       'Each collection should be a `FrozenDict` with string keys.')
+    if rngs is not None and not _is_valid_rngs(rngs):
+      raise ValueError('rngs should be a dictionary mapping strings to `jax.PRNGKey`.')
     new_variables = _unfreeze_variables(variables, mutable)
     with Scope(new_variables, rngs=rngs, mutable=mutable).temporary() as root:
       y = fn(root, *args, **kwargs)
@@ -498,9 +526,52 @@ def init(fn: Callable[..., Any], mutable: CollectionFilter = True) -> Callable[.
     `fn` with the scope partially applied.
   """
   @functools.wraps(fn)
-  def wrapper(rngs, *args, **kwargs):
+  def wrapper(rngs, *args, **kwargs) -> (Any, FrozenVariables):
+    if not _is_valid_rng(rngs) and not _is_valid_rngs(rngs):
+      raise ValueError('First argument passed to an init function should be a `jax.PRNGKey` '
+                       'or a dictionary mapping strings to `jax.PRNGKey`.') 
     if not isinstance(rngs, dict):
-      assert rngs.shape == (2,)
       rngs = {'params': rngs}
     return apply(fn, mutable=mutable)({}, *args, rngs=rngs, **kwargs)
   return wrapper
+
+
+def _is_valid_collection(col: FrozenCollection):
+  if not isinstance(col, FrozenDict):
+    return False
+  for name in col.keys():
+    # any value can be stored in a collection so
+    # only keys can be verified.
+    if not isinstance(name, str):
+      return False
+  return True
+
+
+def _is_valid_variables(variables: FrozenVariables):
+  if not isinstance(variables, (dict, FrozenDict)):
+    return False
+  for name, col in variables.items():
+    if not isinstance(name, str):
+      return False
+    if not _is_valid_collection(col):
+      return False
+  return True
+
+
+def _is_valid_rng(rng: Array):
+  if not isinstance(rng, jnp.ndarray):
+    return False
+  if rng.shape != (2,) or rng.dtype != jnp.uint32:
+    return False
+  return True
+
+
+def _is_valid_rngs(rngs: RNGSequences):
+  if not isinstance(rngs, dict):
+    return False
+  for key, val in rngs.items():
+    if not isinstance(key, str):
+      return False
+    if not _is_valid_rng(val):
+      return False
+  return True


Test patch:
diff --git a/tests/core/lift_test.py b/tests/core/lift_test.py
--- a/tests/core/lift_test.py
+++ b/tests/core/lift_test.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from flax.core import Scope, init, apply, lift
+from flax.core import Scope, init, apply, lift, nn
 
 from jax import random
 from jax import numpy as jnp
@@ -36,6 +36,18 @@ def g(scopes, _):
 
     init(f)(random.PRNGKey(0))
 
+  def test_undefined_param(self):
+    def f(scope):
+      dense = lift.vmap(nn.dense, 
+                        in_axes=(0, None), out_axes=0,
+                        variable_axes={'params': 0},
+                        split_rngs={'params': True})
+      dense(scope.push('dense'), np.ones((3, 2)), 2)
+
+    with self.assertRaisesWithLiteralMatch(ValueError, 'No paramater named "kernel" exists in "/vmap(dense)".'):
+      apply(f)({})
+
+
 
 if __name__ == '__main__':
   absltest.main()
diff --git a/tests/core/scope_test.py b/tests/core/scope_test.py
--- a/tests/core/scope_test.py
+++ b/tests/core/scope_test.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from flax.core import Scope, scope, init, apply, nn
+from flax.core import Scope, scope, freeze, init, apply, nn
 
 from jax import random
 
@@ -64,17 +64,25 @@ def test_inconsistent_param_shapes(self):
     def f(scope):
       scope.param('test', nn.initializers.ones, (4,))
     
-    msg = 'Inconsistent shapes between value and initializer for parameter "test": (2,), (4,)'
+    msg = 'Inconsistent shapes between value and initializer for parameter "test" in "/": (2,), (4,)'
     with self.assertRaisesWithLiteralMatch(ValueError, msg):
-      apply(f)({'params': {'test': np.ones((2,))}})
+      apply(f)(freeze({'params': {'test': np.ones((2,))}}))
 
   def test_mutate_undefined_collection(self):
     def f(scope):
-      scope.put_variable('test', 'test', 123)
+      scope.put_variable('state', 'test', 123)
 
-    with self.assertRaisesWithLiteralMatch(ValueError, 'Collection is not mutable: "test"'):
+    msg = 'Trying to update variable "test" in "/" but collection "state" is immutable.'
+    with self.assertRaisesWithLiteralMatch(ValueError, msg):
       init(f, mutable='params')(random.PRNGKey(0))
 
+  def test_undefined_param(self):
+    def f(scope):
+      nn.dense(scope.push('dense'), np.ones((1, 2)), 2)
+
+    with self.assertRaisesWithLiteralMatch(ValueError, 'No paramater named "kernel" exists in "/dense".'):
+      apply(f)({})
+
 
 if __name__ == '__main__':
   absltest.main()