code: def local_attention2d_spatial_decoder(
      x, kv_dim, heads_dim, feedforward_dim, hparams
  ):
      """Image Transformer decoder with local2D spatial layers."""
      batch_dim, length_dim, model_dim = x.shape.dims
      blocks_h_dim = mtf.Dimension("blocksh", hparams.block_height)
      blocks_w_dim = mtf.Dimension("blocksw", hparams.block_width)
      num_h_blocks_dim = mtf.Dimension(
          "num_h_blocks", hparams.img_len // hparams.block_height
      )
      num_w_blocks_dim = mtf.Dimension(
          "num_w_blocks",
          hparams.img_len * hparams.num_channels // hparams.block_width,
      )
      x = mtf.transpose(
          mtf.reshape(
              x,
              mtf.Shape(
                  [
                      batch_dim,
                      num_h_blocks_dim,
                      blocks_h_dim,
                      num_w_blocks_dim,
                      blocks_w_dim,
                      model_dim,
                  ]
              ),
          ),
          mtf.Shape(
              [
                  batch_dim,
                  num_h_blocks_dim,
                  num_w_blocks_dim,
                  blocks_h_dim,
                  blocks_w_dim,
                  model_dim,
              ]
          ),
      )
      mode = getattr(hparams, "mode", tf_estimator.ModeKeys.TRAIN)
      is_training = mode == tf_estimator.ModeKeys.TRAIN
      # Image Transformer Decoder
      # [ self attention - ffn - residual + dropout] x n
      for layer in range(hparams.num_decoder_layers):
          layer_name = "decoder_layer_%d" % layer
          with tf.variable_scope(layer_name):
              # Self attention layer
              x += layer_prepostprocess_dropout(
                  mtf.layers.local_2d_self_attention_spatial_blocks(
                      mtf.layers.layer_norm(
                          x, model_dim, name="layer_norm_att"
                      ),
                      kv_dim,
                      heads_dim,
                      is_training,
                      memory_h_dim=num_h_blocks_dim,
                      memory_w_dim=num_w_blocks_dim,
                      name="self_att",
                  ),
                  hparams,
              )
              # ffn layer
              x += layer_prepostprocess_dropout(
                  mtf.layers.dense_relu_dense(
                      mtf.layers.layer_norm(
                          x, model_dim, name="layer_norm_ffn"
                      ),
                      feedforward_dim,
                      hparams.dropout,
                      dropout_broadcast_dims=[length_dim],
                  ),
                  hparams,
              )
  
      output = mtf.layers.layer_norm(
          x, model_dim, name="final_layer_norm"
      )
      return output
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import copy
  import mesh_tensorflow
  from tensor2tensor.layers import common_hparams
  from tensor2tensor.layers import common_layers
  from tensor2tensor.utils import mtf_model
  from tensor2tensor.utils import registry
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 3 for calling a global function with example usage: local_attention2d_spatial_decoder(x, kv_dim, heads_dim, feedforward_dim, hparams) and returns: output