class_docstr: A stack of layers for Universal Transformer.
  
  This implementation is largely adapted from t2t universal transformer
  implementation. Reference:
  third_party/py/tensor2tensor/models/research
class_docstr_tok: ["A","stack","of","layers","for","Universal","Transformer",".","This","implementation","is","largely","adapted","from","t2t","universal","transformer","implementation",".","Reference",":","third_party\/py\/tensor2tensor\/models\/research"]
docstr: A highway network layer.
docstr_tok: ["A","highway","network","layer","."]
code: def ut_highway(self, context, layer_inputs, mask):
      """A highway network layer."""
  
      def ut_function(x, step):
          """highway layer implementation."""
          state, inputs, memory = x
          new_state = self.step_preprocess(context, state, step)
          for _ in range(self.num_inrecurrence_layers):
              new_state = self.vanilla_transformer_layer(
                  context, new_state, mask
              )
          transformed_state = new_state
  
          gate_inputs = []
          if "s" in self.gates_inputs:
              gate_inputs.append(state)
          if "t" in self.gates_inputs:
              gate_inputs.append(transformed_state)
          if "i" in self.gates_inputs:
              gate_inputs.append(inputs)
          gate_ffn_layer = self.gate_ffn_layer
  
          transform_gate = self.ffn_layer_multi_inputs(
              context,
              mask,
              gate_inputs,
              ffn_layer_type=gate_ffn_layer,
              activation=mtf.sigmoid,
              preprocess=True,
          )
          if self.couple_carry_transform_gates:
              carry_gate = mtf.sub(1.0, transform_gate, name="carry")
          else:
              carry_gate = self.ffn_layer_multi_inputs(
                  context,
                  mask,
                  gate_inputs,
                  ffn_layer_type=gate_ffn_layer,
                  activation=mtf.sigmoid,
                  preprocess=True,
              )
          new_state = (
              state * carry_gate + transformed_state * transform_gate
          )
  
          mtf.scalar_summary(
              "highway_transform_gate_layer",
              mtf.reduce_mean(transform_gate),
          )
          mtf.scalar_summary(
              "highway_carry_gate_layer", mtf.reduce_mean(carry_gate)
          )
  
          return new_state, inputs, memory
  
      for i in range(self.num_rec_steps):
          layer_inputs = ut_function(layer_inputs, i)
      output, _, _ = layer_inputs
      return output
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import math
  import gin
  import mesh_tensorflow
  from mesh_tensorflow.transformer import transformer
  import tensorflow.compat.v1
label_desc: Use image node_id 12 for calling the UTLayerStack obj's underlying member method code with example usage: obj.ut_highway(context, layer_inputs, mask) and returns: output, new_state,  inputs,  memory