def local_attention_2d(x, hparams, attention_type="local_attention_2d"):
  """Local 2d, self attention layer."""
  # 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
Local 2d, self attention layer.