code: def universal_transformer_highway(
      layer_inputs,
      step,
      hparams,
      ffn_unit,
      attention_unit,
      pad_remover=None,
  ):
      """Universal Transformer with highway connection.
  
  
      It transforms the state using a block containing self-attention and transition
      function  and wrap the whole block with a highway connection.
      (the new state is a combination of the state and the transformed-state
      based on cary/transform gates.)
  
      Interesting observation:
        Controlling the cary/transform gate with the original inputs works usually
        better (i.e. hparams.gates_inputs="i")
  
      Args:
        layer_inputs:
          - state: state
          - inputs: the original embedded inputs (= inputs to the first step)
        step: indicates number of steps taken so far
        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:
        layer_output:
            new_state: new state
            inputs: the original embedded inputs (= inputs to the first step)
  
      """
  
      state, inputs, memory = layer_inputs
      new_state = step_preprocess(state, step, hparams)
  
      for i in range(hparams.num_inrecurrence_layers):
          with tf.variable_scope("rec_layer_%d" % i):
              new_state = ffn_unit(attention_unit(new_state))
  
      transformed_state = new_state
  
      gate_inputs = []
      if "s" in hparams.gates_inputs:
          gate_inputs.append(state)
  
      if "t" in hparams.gates_inputs:
          gate_inputs.append(transformed_state)
  
      if "i" in hparams.gates_inputs:
          gate_inputs.append(inputs)
  
      gate_ffn_layer = hparams.gate_ffn_layer
  
      transform_gate = _ffn_layer_multi_inputs(
          gate_inputs,
          hparams,
          ffn_layer_type=gate_ffn_layer,
          name="transform",
          bias_initializer=tf.constant_initializer(
              hparams.transform_bias_init
          ),
          activation=tf.sigmoid,
          pad_remover=pad_remover,
          preprocess=True,
      )
  
      if hparams.couple_carry_transform_gates:
          carry_gate = tf.subtract(1.0, transform_gate, name="carry")
  
      else:
          carry_gate = _ffn_layer_multi_inputs(
              gate_inputs,
              hparams,
              ffn_layer_type=gate_ffn_layer,
              name="carry",
              bias_initializer=tf.constant_initializer(
                  -hparams.transform_bias_init
              ),
              activation=tf.sigmoid,
              pad_remover=pad_remover,
              preprocess=True,
          )
  
      new_state = (
          state * carry_gate + transformed_state * transform_gate
      )
  
      contrib.summary().scalar(
          "highway_transform_gate_layer", tf.reduce_mean(transform_gate)
      )
  
      contrib.summary().scalar(
          "highway_carry_gate_layer", tf.reduce_mean(carry_gate)
      )
  
      return new_state, inputs, memory
  
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 10 for calling a global function with example usage: universal_transformer_highway(layer_inputs, step, hparams, ffn_unit, attention_unit, pad_remover) and returns: new_state,  inputs,  memory