Improve Error Message: catch invalid initializers earlier?

[relational]

Consider the following snippet, adapted from the flax examples, which has a minor bug ( can you spot it? )

class SimpleDense(nn.Module):
  features: int
  kernel_init: Callable = nn.initializers.lecun_normal
  bias_init: Callable = nn.initializers.zeros

  @nn.compact
  def __call__(self, inputs):
    kernel = self.param('kernel',
                        self.kernel_init,
                        (inputs.shape[-1], self.features))
    bias = self.param('bias', self.bias_init, (self.features,))
                            
    y = lax.dot_general(inputs, kernel,
                        (((inputs.ndim - 1,), (0,)), ((), ())),)
    y = y + bias
    return y

rng = jax.random.PRNGKey(0)
layer = SimpleDense(5)
layer.init(rng, jnp.ones((1, 32, 32, 3))

If we look at the stack trace, it is pretty long and ends at:

/usr/local/lib/python3.7/dist-packages/jax/interpreters/xla.py in apply_primitive(prim, *args, **params)
    227 def apply_primitive(prim, *args, **params):
    228   """Impl rule that compiles and runs a single primitive 'prim' using XLA."""
--> 229   compiled_fun = xla_primitive_callable(prim, *unsafe_map(arg_spec, args), **params)
    230   return compiled_fun(*args)
    231 

/usr/local/lib/python3.7/dist-packages/jax/interpreters/xla.py in arg_spec(x)
    219 
    220 def arg_spec(x):
--> 221   aval = abstractify(x)
    222   try:
    223     return aval, x._device

/usr/local/lib/python3.7/dist-packages/jax/interpreters/xla.py in abstractify(x)
    168   if hasattr(x, '__jax_array__'):
    169     return abstractify(x.__jax_array__())
--> 170   raise TypeError(f"Argument '{x}' of type '{type(x)}' is not a valid JAX type")
    171 
    172 def _make_abstract_python_scalar(typ, val):

TypeError: Argument '<function variance_scaling.<locals>.init at 0x7fe5fbf8bc20>' of type '<class 'function'>' is not a valid JAX type

Reading the stack trace, it sound as if the problem is related to the lax dot product, but the actual problem is that kernel_init should have been set to nn.initializers.lecun_normal() with () at the end ( in contrast with the zeros initializer ??).

Maybe such a problem could be caught earlier in the self.param(...) call, which would help the user pinpoint the problem?

[jheek]

The benevolent dictator orders his subjects to rely on deep stack traces and spurious type assertions rather than readable type errors. If you want good error messages go to the Rust community instead! ;)

On a more serious note:
Unfortunately, we cannot check in Flax that the init function you passed has the correct type signature.
Of course we could pattern match on whether you passed in this specific function that we now is wrong but let's not go down that road just yet :S
I have run into this exact problem many times and I'm sure many others did as well.
The real problem here is that variance_scaling doesn't properly check the types of its input arguments.
So my suggestion is to report an issue (or better a PR) to the JAX project to add some type assertions to nn.initializers.variance_scaling for this common source of errors.

[avital]

Why can't we check that the init function has the right type signature?