Docs: please clarify how to vmap nnx.Module over batch dimension

### System information
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): wsl2 ubuntu 22.04
- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`: 0.8.5, 0.4.30, 0.4.30
- Python version: 3.12 and whatever colab uses today
- GPU/TPU model and memory: nvidia rtx 4090 24gb
- CUDA version (if applicable): 12.5

### Problem you have encountered:

made a nnx.Module to apply to single (1) input output pair example at a time using vmap
the model works on the first pass, but then on the 2nd pass it crashes

### What you expected to happen:

expected vmapping a nnx.Module over a batch of data would not pass some shape on the first call and a different shape on the second call

Thought the docs could help but there's no mention of batch handling in there, seems like the batching for the given mnist example is assumed, which implies autobatching, but that just crashed when i tried it

tried various combinations of nnx jit, jax jit, and they all crashed with various shape bugs
doing one example of mnist at a time with no batching worked, but batching is pretty necessary

dont want to hard code models for some particular batch size, can you please clarify in docs how can we map a nnx.Module over a batch without needing to hardcode the batch size in the kernel dimensions? or, what am i doing wrong in the colab?

side note: i also hit bugs with dropout, the rngs context would crash, deterministic or not
i also wanted to make a dynamic length SGLD scan and early stop when the hessian indicated an optima but jax can't do this apparently :(

### Logs, error messages, etc:


with jax.vmap
with nnx.vmap

### Steps to reproduce:
Whenever possible, please provide a *minimal example*. Please consider submitting it as a Colab link.

NOTE: TO SAVE YOUR TIME, you can skip reading the 2nd codeblock as it's just copypasta to load mnist
https://colab.research.google.com/drive/11gKSydAn3p_fLdQaqw7pkdHmC2zyW12k?usp=sharing