text: Title: Example code for using random_flax_module with flax.linen.BatchNorm (mutable)
  Body: Hi, 
  
  I am trying to use random_flax_module on a class that uses flax.linen.BatchNorm that uses mutable parameters. Is there any example on how to use that? Here is my code:
  
  The model:
  
  ```
  class NSBlock(nn.Module):
  
      train: bool
      dim: int
      ks_1: int = 3
      ks_2: int = 3
      dl_1: int = 1
      dl_2: int = 1
      mp_ks: int = 3
      mp_st: int = 1
  
      @nn.compact
      def __call__(self, x):
          lpad1 = (self.dl_1 * (self.ks_1 - 1)) // 2
          conv_r1 = nn.Conv(features = self.dim,
                  kernel_size = (self.ks_1, self.ks_1),
                  kernel_dilation = self.dl_1,
                  padding= lpad1)
          bn_r1 = nn.BatchNorm(momentum = 0.9, use_running_average = not self.train)
  
          lpad2 = (self.dl_2 * (self.ks_2 - 1)) // 2
          conv_r2 = nn.Conv(features = self.dim,
                  kernel_size = (self.ks_2, self.ks_2),
                  kernel_dilation = self.dl_2,
                  padding= lpad2)
          bn_r2 = nn.BatchNorm(momentum = 0.9, use_running_average = not self.train)
  
          x_0 = conv_r1(x)
          x_1 = bn_r1(x_0)
          x_2 = nn.relu(x_1)
  
          x_3 = conv_r2(x_2)
          x_4 = bn_r2(x_3)
  
          y = x + x_4
          rpad = (self.mp_ks - 1) // 2
          z = nn.max_pool(y, window_shape = (1, self.mp_ks),
                  padding = ((0, 0), (rpad, rpad)), strides = (1, self.mp_st))
          return z
  
  
  class NeuSomaticNet(nn.Module):
  
      train : bool
      dim: int = 64
      nsblocks = [
          [3, 5, 1, 1, 3, 1],
          [3, 5, 1, 1, 3, 2],
          [3, 5, 2, 1, 3, 2],
          [3, 5, 4, 2, 3, 2],
      ]
  
      @nn.compact
      def __call__(self, x_0):
  
          conv1 = nn.Conv(features =  self.dim, kernel_size = (1, 3),
              padding = (0, 1), strides = 1)
          bn1 = nn.BatchNorm(momentum = 0.9, use_running_average = not self.train)
  
          x_1 = conv1(x_0)
          x_2 = bn1(x_1)
          x = nn.relu(x_2)
  
          # x = nn.relu(bn1(conv1(x_0)))
  
          x = nn.max_pool(x, window_shape = (1, 3), padding = ((0, 0), (1, 1)), strides = (1, 1))
  
          for ks_1, ks_2, dl_1, dl_2, mp_ks, mp_st in self.nsblocks:
              rb = NSBlock(self.train, self.dim, ks_1, ks_2, dl_1, dl_2, mp_ks, mp_st)
              x = rb(x)
  
          # Do a transpose before using the dense layer
          x = jnp.transpose(x, (0, 3, 1, 2))
  
          ds = np.prod(list(map(lambda x: x[5], self.nsblocks)))
          fc_dim = self.dim * 32 * 5 // ds
          x2 = jnp.reshape(x, (-1, fc_dim))
  
          x3 = nn.Dense(features = 240)(x2)
          o1 = nn.Dense(features = 4)(x3)
  
          return o1
  
  ```
  
  The numpyro code:
  
  ```
  def neu_model(X, Y = None, D_Y = 1, train = True):
      N, H, W, C = X.shape
      lclass = 4
  
      module = NeuSomaticNet(train = train)
  
      key1, key2 = random.split(random.PRNGKey(0), 2)
  
      vars_in = module.init(key1, X)
      lvals, mutated_vars = module.apply(vars_in, X, mutable=['batch_stats'])
  
      net = random_flax_module("net", module, prior = (
                              lambda name, shape:
                                  dist.Cauchy() if "bias" in name
                                  else dist.Normal(0, 1)),
                              input_shape = (N, H, W, C), mutable=['batch_stats'])
  
      lvals = net(X)
      with numpyro.plate("batch", X.shape[0]):
          numpyro.sample("Y", dist.Categorical(logits = lvals), obs = Y)
  
  ```
  
  Part of the inference code:
  
  ```
  def run_inference_DELTA(model_l, rng_key, X, Y, num_samples, svi_steps):
      print("Starting DELTA inference")
      guide = AutoDelta(model_l)
      svi = SVI(model_l, guide, optim.Adam(0.001), Trace_ELBO())
      svi_result = svi.run(random.PRNGKey(100), svi_steps, X, Y, train = True)
      # create_loss_plot(args, svi_result.losses)
      print("AutoDelta inference done")
      params = svi_result.params
      samples_1 = guide.sample_posterior(
          random.PRNGKey(1), params, sample_shape=(num_samples,)
      )
      return samples_1
  
  rng_key, rng_key_predict = random.split(random.PRNGKey(10))
  
   samples_1 = run_inference_DELTA(neu_model, rng_key, inputs_f, labels_f, num_samples, svi_steps)
  
  ```
  
  Here is the error that I am getting:
  
  ```
  
  Starting DELTA inference
  ERROR 2022-07-04 11:55:04,285 __main__             Traceback (most recent call last):
    File "/home/bandyopn/global/code/neusomatic_uq/neusomatic/python/train_uq.py", line 797, in <module>
      checkpoint = train_neusomatic(args.candidates_tsv, args.validation_candidates_tsv,
    File "/home/bandyopn/global/code/neusomatic_uq/neusomatic/python/train_uq.py", line 633, in train_neusomatic
      samples_1 = run_inference_DELTA(neu_model, rng_key, inputs_f, labels_f, num_samples, svi_steps)
    File "/home/bandyopn/global/code/neusomatic_uq/neusomatic/python/train_uq.py", line 673, in run_inference_DELTA
      svi_result = svi.run(random.PRNGKey(100), svi_steps, X, Y, train = True)
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/infer/svi.py", line 342, in run
      svi_state = self.init(rng_key, *args, **kwargs)
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/infer/svi.py", line 180, in init
      guide_trace = trace(guide_init).get_trace(*args, **kwargs, **self.static_kwargs)
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/handlers.py", line 171, in get_trace
      self(*args, **kwargs)
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/primitives.py", line 105, in __call__
      return self.fn(*args, **kwargs)
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/primitives.py", line 105, in __call__
      return self.fn(*args, **kwargs)
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/infer/autoguide.py", line 430, in __call__
      self._setup_prototype(*args, **kwargs)
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/infer/autoguide.py", line 404, in _setup_prototype
      super()._setup_prototype(*args, **kwargs)
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/infer/autoguide.py", line 154, in _setup_prototype
      ) = initialize_model(
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/infer/util.py", line 654, in initialize_model
      (init_params, pe, grad), is_valid = find_valid_initial_params(
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/infer/util.py", line 395, in find_valid_initial_params
      (init_params, pe, z_grad), is_valid = _find_valid_params(
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/infer/util.py", line 381, in _find_valid_params
      _, _, (init_params, pe, z_grad), is_valid = init_state = body_fn(init_state)
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/infer/util.py", line 331, in body_fn
      model_trace = trace(seeded_model).get_trace(*model_args, **model_kwargs)
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/handlers.py", line 171, in get_trace
      self(*args, **kwargs)
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/primitives.py", line 105, in __call__
      return self.fn(*args, **kwargs)
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/primitives.py", line 105, in __call__
      return self.fn(*args, **kwargs)
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/primitives.py", line 105, in __call__
      return self.fn(*args, **kwargs)
    [Previous line repeated 2 more times]
    File "/home/bandyopn/global/code/neusomatic_uq/neusomatic/python/train_uq.py", line 656, in neu_model
      net = random_flax_module("net", module, prior = (
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/contrib/module.py", line 355, in random_flax_module
      nn = flax_module(
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/contrib/module.py", line 79, in flax_module
      assert (nn_state is None) == (nn_params is None)
  AssertionError
  
  
  Traceback (most recent call last):
    File "/home/bandyopn/global/code/neusomatic_uq/neusomatic/python/train_uq.py", line 821, in <module>
      raise e
    File "/home/bandyopn/global/code/neusomatic_uq/neusomatic/python/train_uq.py", line 797, in <module>
      checkpoint = train_neusomatic(args.candidates_tsv, args.validation_candidates_tsv,
    File "/home/bandyopn/global/code/neusomatic_uq/neusomatic/python/train_uq.py", line 633, in train_neusomatic
      samples_1 = run_inference_DELTA(neu_model, rng_key, inputs_f, labels_f, num_samples, svi_steps)
    File "/home/bandyopn/global/code/neusomatic_uq/neusomatic/python/train_uq.py", line 673, in run_inference_DELTA
      svi_result = svi.run(random.PRNGKey(100), svi_steps, X, Y, train = True)
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/infer/svi.py", line 342, in run
      svi_state = self.init(rng_key, *args, **kwargs)
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/infer/svi.py", line 180, in init
      guide_trace = trace(guide_init).get_trace(*args, **kwargs, **self.static_kwargs)
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/handlers.py", line 171, in get_trace
      self(*args, **kwargs)
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/primitives.py", line 105, in __call__
      return self.fn(*args, **kwargs)
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/primitives.py", line 105, in __call__
      return self.fn(*args, **kwargs)
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/infer/autoguide.py", line 430, in __call__
      self._setup_prototype(*args, **kwargs)
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/infer/autoguide.py", line 404, in _setup_prototype
      super()._setup_prototype(*args, **kwargs)
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/infer/autoguide.py", line 154, in _setup_prototype
      ) = initialize_model(
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/infer/util.py", line 654, in initialize_model
      (init_params, pe, grad), is_valid = find_valid_initial_params(
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/infer/util.py", line 395, in find_valid_initial_params
      (init_params, pe, z_grad), is_valid = _find_valid_params(
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/infer/util.py", line 381, in _find_valid_params
      _, _, (init_params, pe, z_grad), is_valid = init_state = body_fn(init_state)
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/infer/util.py", line 331, in body_fn
      model_trace = trace(seeded_model).get_trace(*model_args, **model_kwargs)
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/handlers.py", line 171, in get_trace
      self(*args, **kwargs)
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/primitives.py", line 105, in __call__
      return self.fn(*args, **kwargs)
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/primitives.py", line 105, in __call__
      return self.fn(*args, **kwargs)
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/primitives.py", line 105, in __call__
      return self.fn(*args, **kwargs)
    [Previous line repeated 2 more times]
    File "/home/bandyopn/global/code/neusomatic_uq/neusomatic/python/train_uq.py", line 656, in neu_model
      net = random_flax_module("net", module, prior = (
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/contrib/module.py", line 355, in random_flax_module
      nn = flax_module(
    File "/home/bandyopn/.local/lib/python3.8/site-packages/numpyro/contrib/module.py", line 79, in flax_module
      assert (nn_state is None) == (nn_params is None)
  AssertionError
  
  ```
  
  Could anyone please help me, regarding how to use random_flax_module with BatchNorm? Thanks.