Nested `nn.vmap` of `nn.Module` with `nn.BatchNorm` causes unexpected ValueError

Dear flax team,

this is probably not a bug, but rather a user error. As a relatively new user of flax (migrating from haiku), I found the interaction of nested vmaps with batch normalization quite puzzling. After reading the documentation and related issues, I still couldn't figure out my mistake. Thank you for your help!

### System information
- OS Platform and Distribution: `Ubuntu 20.04`
- Flax, jax, jaxlib versions: `flax: 0.7.2 | jax: 0.4.14 | jaxlib: 0.4.14+cuda11.cudnn86`
- Python version: `3.11.4`
- GPU model: `NVIDIA A100-SXM4-40GB`

### Problem you have encountered:
Transforming a (inner) `nn.Module` that uses `nn.BatchNorm` by applying nested `nn.vmap` with `"variable_axes": {"params": None, "batch_stats": None}` produces
```
ValueError: vmap has mapped output but out_axes is None
```
even though `out_axes` and `axis_name` are provided to both `nn.vmap` and `nn.BatchNorm`. To be precise, this error occurs when applying the twice vmapped inner module. I provide a MWE below.

Notably, the nesting works flawlessly when the inner module does not use batch normalization.

### What you expected to happen:
During training, I expect the batch normalization statistics to be computed across all axes specified in the nested vmap structure. During inference, I expect the same running averages to be used in all calls to the inner module, specified again by the nested vmap structure.

If only a single `nn.vmap` is used, this is exactly what happens. You can see it in the MWE below.

As stated in the [documentation](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/_autosummary/flax.linen.vmap.html), `nn.vmap` returns a module with the same arguments as its target. Therefore, I expect to be able to apply a similar `nn.vmap` transformation again, adjusting for changes in `in_axes` and `out_axes` if necessary, even if batch normalization is used.

### Steps to reproduce:

In this MWE, I process a minibatch with two batch dimensions by applying `nn.vmap` twice to a module. I want to use this setup during training, so I set `train = True` and get a ValueError. There is no error if `train = False`.

```python
import jax.numpy as jnp
import flax.linen as nn
from jax.random import PRNGKey
import jax

class Foo(nn.Module):
features: int
axis_name: str

@nn.compact
def __call__(self, train, x):
x = nn.Dense(features=self.features)(x)
x = nn.BatchNorm(axis_name=self.axis_name)(x, use_running_average=not train)
return x

# prepare data
BATCH_SIZE_0 = 2
BATCH_SIZE_1 = 10
INPUT_SIZE = 3
FEATURE_SIZE = 5

minibatch = jnp.ones((BATCH_SIZE_0, BATCH_SIZE_1, INPUT_SIZE))

# prepare modules and instances
vmap_config = {
"variable_axes": {"params": None, "batch_stats": None},
"split_rngs": {"params": False, "batch_stats": False},
"in_axes": (None, 0),
"out_axes": 0,
"axis_name": "batch",
}

VmapFoo = nn.vmap(Foo, **vmap_config)
VmapVmapFoo = nn.vmap(VmapFoo, **vmap_config)

foo = Foo(features=FEATURE_SIZE, axis_name=None)
vmap_foo = VmapFoo(features=FEATURE_SIZE, axis_name="batch")
vmap_vmap_foo = VmapVmapFoo(features=FEATURE_SIZE, axis_name="batch")

# get variables from foo
rngs = {"params": PRNGKey(0), "batch_stats": PRNGKey(1)}
variables = foo.init(rngs, False, minibatch[0][0])

# apply modules
train = True
mutable = ["batch_stats"]

foo_output, foo_state = foo.apply(variables, train, minibatch[0][0], mutable=mutable)
assert foo_output.shape == (FEATURE_SIZE,)

vmap_foo_output, vmap_foo_state = vmap_foo.apply(variables, train, minibatch[0], mutable=mutable)
assert vmap_foo_output.shape == (BATCH_SIZE_1, FEATURE_SIZE)
# jax.tree_map(lambda f, vf: (f.shape == vf.shape), foo_state, vmap_foo_state) is True at all leaves

vmap_vmap_foo_output, vmap_vmap_foo_state = vmap_vmap_foo.apply(variables, train, minibatch, mutable=mutable)
assert vmap_vmap_foo_output.shape == (BATCH_SIZE_0, BATCH_SIZE_1, FEATURE_SIZE)
```

This returns a ValueError, even though `out_axes` and `axis_name` are specified in `vmap_config`:
```
Traceback (most recent call last):
File "", line 57, in
vmap_vmap_foo_output, vmap_vmap_foo_state = vmap_vmap_foo.apply(variables, train, minibatch, mutable=mutable)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: vmap has mapped output (axis_name='batch') but out_axes is None
```

Understanding my mistake in this setup will surely improve my understanding of state handling in flax. Thank you for your guidance!