AutoRegressive Decoding currently fails if input prompt > 1

Provide as much information as possible. At least, this should include a description of your issue and steps to reproduce the problem. If possible also provide a summary of what steps or workarounds you have already tried.

### Problem you have encountered:

I want to run a model based on `flax.linen.SelfAttention` in auto-regressive mode and pass an input prompt > 1. This however does not seem possible at the moment, *e.g.*:

```python
import jax
import jax.numpy as jnp
from flax.linen import SelfAttention

attn_layer = SelfAttention(1, decode=True, use_bias=False)

batch_size = 1
max_decoder_length = 4
hidden_size = 2
prompt_length = 2 # setting this to 1 would work

init_variables = attn_layer.init(jax.random.PRNGKey(0), jnp.ones((batch_size, max_decoder_length, hidden_size)), deterministic=True)

params = init_variables["params"]
cache = init_variables["cache"]

dummy_prompt = jnp.arange(batch_size * prompt_length * hidden_size).reshape((batch_size, prompt_length, hidden_size))

output, cache = attn_layer.apply({"params": params, "cache": cache}, dummy_prompt, mutable=["cache"], deterministic=True)
```

leads to an error. Also check [this](https://colab.research.google.com/drive/1vBTsCiWvdK7X83SXrAtfdU24DR_QY8R8?usp=sharing) notebook.

### What you expected to happen:

Instead, the code should work and the first `len(prompt_length)` cache variables should be stored.

### Logs, error messages, etc:

```
~/python_bin/flax/linen/module.py in wrapped_module_method(*args, **kwargs)
273 _context.module_stack.append(self)
274 try:
--> 275 y = fun(self, *args, **kwargs)
276 if _context.capture_stack:
277 filter_fn = _context.capture_stack[-1]

~/python_bin/flax/linen/attention.py in __call__(self, inputs_q, inputs_kv, mask, deterministic)
265 expected_shape = tuple(batch_dims) + (1, num_heads, depth_per_head)
266 if expected_shape != query.shape:
--> 267 raise ValueError('Autoregressive cache shape error, '
268 'expected query shape %s instead got %s.' %
269 (expected_shape, query.shape))

ValueError: Autoregressive cache shape error, expected query shape (1, 1, 1, 2) instead got (1, 2, 1, 2).
```

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

See code/colab above