code: def transformer_revnet_decoder(
      decoder_input,
      encoder_output,
      decoder_self_attention_bias,
      encoder_decoder_attention_bias,
      hparams,
      name="decoder",
  ):
      """A stack of transformer layers.
  
      Args:
        decoder_input: a Tensor
        encoder_output: a Tensor
        decoder_self_attention_bias: bias Tensor for self-attention
          (see common_attention.attention_bias())
        encoder_decoder_attention_bias: bias Tensor for encoder-decoder attention
          (see common_attention.attention_bias())
        hparams: hyperparameters for model
        name: a string
  
      Returns:
        y: a Tensors
      """
  
      def f(x, side_input):
          """f(x) for reversible layer, self-attention and enc-dec attention."""
          decoder_self_attention_bias = side_input[0]
          encoder_decoder_attention_bias = side_input[1]
          encoder_output = side_input[2]
  
          old_hid_size = hparams.hidden_size
          hparams.hidden_size = old_hid_size // 2
  
          with tf.variable_scope("self_attention"):
              y = common_attention.multihead_attention(
                  common_layers.layer_preprocess(x, hparams),
                  None,
                  decoder_self_attention_bias,
                  hparams.attention_key_channels or hparams.hidden_size,
                  hparams.attention_value_channels
                  or hparams.hidden_size,
                  hparams.hidden_size,
                  hparams.num_heads,
                  hparams.attention_dropout,
              )
              y = common_layers.layer_postprocess(x, y, hparams)
              if encoder_output is not None:
                  with tf.variable_scope("encdec_attention"):
                      y = common_attention.multihead_attention(
                          common_layers.layer_preprocess(x, hparams),
                          encoder_output,
                          encoder_decoder_attention_bias,
                          hparams.attention_key_channels
                          or hparams.hidden_size,
                          hparams.attention_value_channels
                          or hparams.hidden_size,
                          hparams.hidden_size,
                          hparams.num_heads,
                          hparams.attention_dropout,
                      )
                      y = common_layers.layer_postprocess(x, y, hparams)
          hparams.hidden_size = old_hid_size
          return y
  
      def g(x):
          """g(x) for reversible layer, feed-forward layer."""
          old_hid_size = hparams.hidden_size
          hparams.hidden_size = old_hid_size // 2
          with tf.variable_scope("ffn"):
              y = transformer.transformer_ffn_layer(
                  common_layers.layer_preprocess(x, hparams), hparams
              )
              y = common_layers.layer_postprocess(x, y, hparams)
          hparams.hidden_size = old_hid_size
          return y
  
      x1, x2 = tf.split(decoder_input, 2, axis=-1)
  
      with tf.variable_scope(name):
          y1, y2 = contrib.layers().rev_block(
              x1,
              x2,
              f,
              g,
              num_layers=hparams.num_hidden_layers,
              f_side_input=[
                  decoder_self_attention_bias,
                  encoder_decoder_attention_bias,
                  encoder_output,
              ],
              is_training=hparams.mode == tf_estimator.ModeKeys.TRAIN,
          )
          y = tf.concat([y1, y2], axis=-1)
          return common_layers.layer_preprocess(y, hparams)
  
file_path: tensor2tensor/tensor2tensor/models/research/transformer_revnet.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from tensor2tensor.layers import common_attention
  from tensor2tensor.layers import common_layers
  from tensor2tensor.models import transformer
  from tensor2tensor.utils import contrib
  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: transformer_revnet_decoder(decoder_input, encoder_output, decoder__attention_bias, encoder_decoder_attention_bias, hparams, name) and returns: y, y, common_layers