code: def get_ut_layer(
      x, hparams, ffn_unit, attention_unit, pad_remover=None
  ):
      """Provides the function that is used in universal transforemr steps.
  
      Args:
        x: input
        hparams: model hyper-parameters
        ffn_unit: feed-forward unit
        attention_unit: multi-head attention unit
        pad_remover: to mask out padding in convolutional layers (efficiency).
  
      Returns:
        ut_function and the ut_initializer
  
      Raises:
        ValueError: Unknown recurrence type
      """
  
      if hparams.recurrence_type == "basic":
          ut_initializer = (x, x, x)  # (state, input, memory)
          ut_function = functools.partial(
              universal_transformer_basic,
              hparams=hparams,
              ffn_unit=ffn_unit,
              attention_unit=attention_unit,
          )
  
      elif hparams.recurrence_type == "highway":
          ut_initializer = (x, x, x)  # (state, input, memory)
          ut_function = functools.partial(
              universal_transformer_highway,
              hparams=hparams,
              ffn_unit=ffn_unit,
              attention_unit=attention_unit,
              pad_remover=pad_remover,
          )
  
      elif hparams.recurrence_type == "skip":
          ut_initializer = (x, x, x)  # (state, input, memory)
          ut_function = functools.partial(
              universal_transformer_skip,
              hparams=hparams,
              ffn_unit=ffn_unit,
              attention_unit=attention_unit,
              pad_remover=pad_remover,
          )
  
      elif hparams.recurrence_type == "dwa":
          # memory contains the original input + all the states
          memory_size = hparams.num_rec_steps + 1
  
          # prepare initializer:
          memory_empty = tf.zeros(
              [memory_size] + common_layers.shape_list(x)
          )
  
          # filling the first slot with the original input
          memory = fill_memory_slot(memory_empty, x, 0)
  
          ut_initializer = (x, x, memory)  # (state, input, memory)
          ut_function = functools.partial(
              universal_transformer_depthwise_attention,
              hparams=hparams,
              ffn_unit=ffn_unit,
              attention_unit=attention_unit,
          )
  
      elif hparams.recurrence_type == "gru":
          ut_initializer = (x, x, x)  # (state, input, memory)
          ut_function = functools.partial(
              universal_transformer_with_gru_as_transition_function,
              hparams=hparams,
              ffn_unit=ffn_unit,
              attention_unit=attention_unit,
              pad_remover=pad_remover,
          )
  
      elif hparams.recurrence_type == "lstm":
          memory = tf.zeros(common_layers.shape_list(x))
          ut_initializer = (x, x, memory)  # (state, input, memory)
          ut_function = functools.partial(
              universal_transformer_with_lstm_as_transition_function,
              hparams=hparams,
              ffn_unit=ffn_unit,
              attention_unit=attention_unit,
              pad_remover=pad_remover,
          )
  
      else:
          raise ValueError(
              "Unknown recurrence type: %s" % hparams.recurrence_type
          )
  
      return ut_function, ut_initializer
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import copy
  import functools
  from six.moves import range
  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 expert_utils
  import tensorflow.compat.v1
label_desc: Use image node_id 4 for calling a global function with example usage: get_ut_layer(x, hparams, ffn_unit, attention_unit, pad_remover) and returns: ut_function,  ut_initializer