code: 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
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  from six.moves import range
  from tensor2tensor.layers import common_attention
  from tensor2tensor.layers import common_layers
  from tensor2tensor.utils import expert_utils
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 2 for calling a global function with example usage: local_attention_2d(x, hparams, attention_type) and returns: y