text: Title: Allow prior to be a callable in random module
  Body: Currently, [random_flax_module](http://num.pyro.ai/en/stable/primitives.html#numpyro.contrib.module.random_flax_module) and [random_haiku_module](http://num.pyro.ai/en/stable/primitives.html#numpyro.contrib.module.random_haiku_module) assume that `prior` is either a distribution or a dict of distributions. In case we have a large network and want to specify priors for specific parameters like the weight of the last layer, we need to find corresponding names of those parameters. It is a bit inconvenient. It would be nice to allow a callable `prior` argument like
  ```
  def get_prior(name, shape):
      if "weight_last" in name:
          return dist.Normal(0, 1).expand([shape])
  ```
  where `name` is the parameter full name and `shape` is its corresponding shape.
  
  This is a nice feature and supporting this is not complicated. We just need to add a [check](https://github.com/pyro-ppl/numpyro/blob/d3a4ca116d8d99372fae030fe04990db6a522055/numpyro/contrib/module.py#L226) here such that when `prior` is a callable, we get `d` by calling `prior(flatten_name, param_shape)`. We'll need a few tests for this feature and this would be a nice chance to get familiar with using flax/haiku in numpyro.