code: def compress_self_attention_layer(x, hparams, name=None):
    
    with tf.variable_scope(name, default_name="compress_self_attention"):
      x, xshape, _ = cia.maybe_reshape_4d_to_3d(x)
      y = common_attention.multihead_attention(
          common_layers.layer_preprocess(x, hparams),
          None,
          None,
          hparams.attention_key_channels or hparams.hidden_size,
          hparams.attention_value_channels or hparams.hidden_size,
          hparams.hidden_size, hparams.num_heads,
          hparams.attention_dropout)
      res = common_layers.layer_postprocess(x, y, hparams)
      return tf.reshape(res, xshape)
docs: Attend function.