code: def local_attention1d_masked_decoder(x, kv_dim, heads_dim,
                                       feedforward_dim, hparams):
    
    print(x)
    _, length_dim, model_dim = x.shape.dims
    for layer in range(hparams.num_decoder_layers):
      layer_name = "decoder_layer_%d" % layer
      with tf.variable_scope(layer_name):
        # Self attention layer
        length_per_split = mtf.tensor_dim_to_size_per_split(
            hparams.layout, hparams.mesh_shape, length_dim)
        x += layer_prepostprocess_dropout(
            mtf.layers.masked_local_attention_1d(
                mtf.layers.layer_norm(x, model_dim, name="layer_norm_att"),
                kv_dim,
                heads_dim,
                window_size=hparams.block_length,
                length_per_split=length_per_split,
                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
docs: Image Transformer decoder with local1D masked layers.