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: Implementation of various gating layers.
docstr_tok: ["Implementation","of","various","gating","layers","."]
code: def gating(self, context, x, transformed_x, mask):
      """Implementation of various gating layers."""
      gate_ffn_layer = self.gate_ffn_layer
      if self.gating_type == "highway":
          gate_inputs = [x]
          transform_gate = self.ffn_layer_multi_inputs(
              context,
              mask,
              gate_inputs,
              ffn_layer_type=gate_ffn_layer,
              activation=mtf.sigmoid,
              preprocess=True,
          )
          carry_gate = self.ffn_layer_multi_inputs(
              context,
              mask,
              gate_inputs,
              ffn_layer_type=gate_ffn_layer,
              activation=mtf.sigmoid,
              preprocess=True,
          )
          new_state = x * carry_gate + transformed_x * transform_gate
          return new_state
      elif self.gating_type == "gru":
          gate_inputs = [x, transformed_x]
          transition_function_update_gate = self.ffn_layer_multi_inputs(
              context,
              mask,
              gate_inputs,
              ffn_layer_type=gate_ffn_layer,
              activation=mtf.sigmoid,
              preprocess=True,
          )
          transition_function_reset_gate = self.ffn_layer_multi_inputs(
              context,
              mask,
              gate_inputs,
              ffn_layer_type=gate_ffn_layer,
              activation=mtf.sigmoid,
              preprocess=True,
          )
  
          reset_state = transition_function_reset_gate * x
          gate_inputs = [reset_state, transformed_x]
          transition_function_candidate = self.ffn_layer_multi_inputs(
              context,
              mask,
              gate_inputs,
              ffn_layer_type=gate_ffn_layer,
              activation=mtf.sigmoid,
              preprocess=True,
          )
  
          transition_function_output = (
              (1 - transition_function_update_gate) * transformed_x
              + transition_function_update_gate
              * transition_function_candidate
          )
          return transition_function_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 8 for calling the UTLayerStack obj's underlying member method code with example usage: obj.gating(context, x, transformed_x, mask) and returns: new_state, transition_function_output