Error occurs in `nn.vmap` while `variable_axes` is a nested dict

Hello, I am trying to use `nn.vmap` to vectorize a module in flax, and some parameters are shared over batches. Here is an example:
```
class Test(nn.Module):

@nn.compact
def __call__(self, x):

x = nn.Dense(4)(x)
return nn.Dense(2)(x)

BatchTest = nn.vmap(
Test,
in_axes=0,
out_axes=0,
variable_axes={'params': {'Dense_0': {'bias': None, 'kernel': None}, 'Dense_1': {'bias': 0, 'kernel': 0}}},
split_rngs={'params': False},
)

_params = BatchTest().init({'params': jax.random.PRNGKey(42)}, jnp.zeros((3, 3)))
```
Here the first dense layer's parameters are expected to be shared over batches. However, I got the error report:
```
---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
Cell In[13], line 9
1 BatchTest = nn.vmap(
2 Test,
3 in_axes=0,
(...)
6 split_rngs={'params': False},
7 )
----> 9 _params = BatchTest().init({'params': jax.random.PRNGKey(42)}, jnp.zeros((3, 3)))
[... skipping hidden 11 frame]
File ~/miniconda3/envs/jax/lib/python3.9/site-packages/jax/_src/tree_util.py:243, in (.0)
210 """Maps a multi-input function over pytree args to produce a new pytree.
211
212 Args:
(...)
240 [[5, 7, 9], [6, 1, 2]]
241 """
242 leaves, treedef = tree_flatten(tree, is_leaf)
--> 243 all_leaves = [leaves] + [treedef.flatten_up_to(r) for r in rest]
244 return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
ValueError: Expected dict, got ({},).
```
I wonder if there are ways to solve the problem?