Perceiver: identifying pre-trained key and query weights

The [Perceiver model Colabs](https://github.com/deepmind/deepmind-research/tree/master/perceiver#colabs) are an excellent resource -- thank you for providing them.

I'm trying to correctly identify the key and query weight matrices for analysis. However, I am a bit confused about the internal naming conventions used and some of the model meta-parameters.

Take [colabs/masked_language_modelling.ipynb](https://colab.research.google.com/github/deepmind/deepmind_research/blob/master/perceiver/colabs/masked_language_modelling.ipynb) as an example. My goal is to identify the pre-trained key and query weights (and biases) such that I can manually simulate and analyse the transformation from a vector of the latent variables to its transformation into keys and queries before and after each module of self-attention (I would then like to do the same for cross-attention).

Running this workbook through, we download the pre-trained model parameters (`language_perceiver_io_bytes.pickle`) and put it into the dictionary `params` (which, as values, stores dictionaries).

`params` stores 231 key-value pairs. Some string keys share the same starting characters, e.g.,
```
'perceiver_encoder/~/self_attention_1/attention/linear_1',
'perceiver_encoder/~/self_attention_1/attention/linear_3',
'perceiver_encoder/~/self_attention_1/attention/linear',
'perceiver_encoder/~/self_attention_1/attention/linear_2',
'perceiver_encoder/~/self_attention_1/layer_norm',
'perceiver_encoder/~/self_attention_1/layer_norm_1',
'perceiver_encoder/~/self_attention_1/mlp/linear',
'perceiver_encoder/~/self_attention_1/mlp/linear_1'
```
Each of the `linear` elements above have are dictionaries with two string keys: `b` (I believe for bias) and `w` (I believe for weights). However, which, of the four `/attention/linear` parameters are being used to (from a 256-sized vector, the latent vector) create the keys, queries, and values used in the self-attention calculation?

Using
```
print("linear b shape:",np.shape(params['perceiver_encoder/~/self_attention_1/attention/linear']['b']))
print("linear w shape:",np.shape(params['perceiver_encoder/~/self_attention_1/attention/linear']['w']))
print("linear1 b shape:",np.shape(params['perceiver_encoder/~/self_attention_1/attention/linear_1']['b']))
print("linear1 w shape:",np.shape(params['perceiver_encoder/~/self_attention_1/attention/linear_1']['w']))
print("linear2 b shape:",np.shape(params['perceiver_encoder/~/self_attention_1/attention/linear_2']['b']))
print("linear2 w shape:",np.shape(params['perceiver_encoder/~/self_attention_1/attention/linear_2']['w']))
print("linear3 b shape:",np.shape(params['perceiver_encoder/~/self_attention_1/attention/linear_3']['b']))
print("linear3 w shape:",np.shape(params['perceiver_encoder/~/self_attention_1/attention/linear_3']['w']))
```
I can see the shapes are
```
linear b shape: (256,)
linear w shape: (1280, 256)
linear1 b shape: (256,)
linear1 w shape: (1280, 256)
linear2 b shape: (1280,)
linear2 w shape: (1280, 1280)
linear3 b shape: (1280,)
linear3 w shape: (1280, 1280)
```

In this case, I've also noticed that there ought to be 8 self-attention heads, 8*32=256 qk channels, 1280 v channels (for the encoder), and 768 v channels (for the decoder).

What I'm confused about is that I thought the number of latent variables remained fixed (in this case, 32 per self-attention head, or 256 in total). Yet, I see no weight matrix of size (256, 256), which I would expect for self-attention. Instead I see two (per attention module) weight matrices which can accept a vector of size 256 -- shaped (1280, 256) -- and then two which accept a vector of size 1280 -- shaped (1280, 1280). Due to their shapes, both of these will output a vector of size 1280 (which in the first two cases doesn't comport with the number of 'biases' found in `b`, since there are only 256 of them).

Apologies if I've made some very obvious oversight or basic error.