code: def local_attention_2d(x, hparams, attention_type="local_attention_2d"):
    
    # self-attention
    with tf.variable_scope("local_2d_self_att"):
      y = common_attention.multihead_attention_2d(
          x,
          None,
          hparams.attention_key_channels or hparams.hidden_size,
          hparams.attention_value_channels or hparams.hidden_size,
          hparams.hidden_size,
          hparams.num_heads,
          attention_type=attention_type,
          query_shape=hparams.query_shape,
          memory_flange=hparams.memory_flange,
          name="self_attention")
    return y
docs: Local 2d, self attention layer.