Scan with split parameters

### Description of the model to be implemented

If I have 6 layers stacked of a model (think transformer style), I would like to have their setup be batched by Flax linen. So for example in `setup` I could call.

```python
self.layers = [SeqInternal() for _ in range(self.n_layers)]
```

However it seems more natural to do this with scan.

```python
SeqInternalStack = nn.scan(SeqInternal,
split_rngs={"params" : True, "dropout": True},
in_axes=0,
out_axes=0,
variable_axes={"params": 0},
length=self.n_layers)
self.layers = SeqInternalStack()
```

In theory (?) this would batch the setup call for each of my SeqInternal layers.

However when I try to do this, things are really slow. Is there an issue with spliting on params for scan? I don't see it in any examples.