code: def local_attention1d_spatial_decoder(
      x, kv_dim, heads_dim, feedforward_dim, hparams
  ):
      """Image Transformer decoder with local1D spatial layers."""
      batch_dim, length_dim, model_dim = x.shape.dims
      blocks_w_dim = mtf.Dimension("blocksw", hparams.block_length)
      num_w_blocks_dim = mtf.Dimension(
          "num_wblocks", length_dim.size // blocks_w_dim.size
      )
      x = mtf.reshape(
          x,
          mtf.Shape(
              [batch_dim, num_w_blocks_dim, blocks_w_dim, model_dim]
          ),
      )
      # [ self attention - ffn - residual + dropout] x n
      mode = getattr(hparams, "mode", tf_estimator.ModeKeys.TRAIN)
      is_training = mode == tf_estimator.ModeKeys.TRAIN
      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_self_attention_spatial_blocks(
                      mtf.layers.layer_norm(
                          x, model_dim, name="layer_norm_att"
                      ),
                      kv_dim,
                      heads_dim,
                      is_training,
                      memory_w_dim=blocks_w_dim,
                      mask_right=True,
                      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,
                      is_training,
                      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 2 for calling a global function with example usage: local_attention1d_spatial_decoder(x, kv_dim, heads_dim, feedforward_dim, hparams) and returns: output