def project(hidden_states, proj_layer):
      """projects hidden states correctly to key/query states"""
      # self-attn
      # (batch_size, n_heads, seq_length, dim_per_head)
      return shape(proj_layer(hidden_states))
projects hidden states correctly to key/query states
extra/models/t5.py