def highway(input_emb, output_size, dropout_ratio=None, mode=None,
            layer_norm=True, carry_bias=1.0):
  """A single highway layer as described here: https://arxiv.org/abs/1505.00387.

  This is used as a generalization of a single hidden layer in a standard
  feedforward network. The difference is that the hidden layer with the
  non-linear activation is interpolated with a linear projection, enabling a
  "highway" for gradients.

  Args:
    input_emb: tensor<float> [..., embedding_size]
    output_size: An integer indicating the size of the hidden layer.
    dropout_ratio: The probability of dropping out each unit in the activation.
        This can be None, and is only applied during training.
    mode: One of the keys from tf.estimator.ModeKeys.
    layer_norm: Boolean indicating whether we should apply layer normalization.
    carry_bias: Float indicating the initial bias for the carry gate. This is
        typically set to a positive number to encourage information flow through
        the "highway" initially. This is analogous to the forget bias typically
        found in LSTMs.

  Returns:
    output_emb: A Tensor with the same shape as `input_emb`, except for the last
        dimension which will have size `output_size instead.
  """
  if mode == tf_estimator.ModeKeys.TRAIN and dropout_ratio is not None:
    input_emb = tf.nn.dropout(input_emb, 1.0 - dropout_ratio)

  with tf.variable_scope("joint_linear_layer"):
    # We need to compute three projections from the input embeddings for a
    # highway layer: (1) the hidden layer with a non-linear activation, (2) the
    # transform gate with a sigmoid activation, and (3) the projected input with
    # which the hidden layer is interpolated. These projections are combined for
    # better parallelism.
    joint_linear = tf.layers.dense(input_emb, output_size * 3)
  hidden_emb, carry_gate, project_input_emb = tf.split(joint_linear, 3, -1)
  hidden_emb = tf.nn.relu(hidden_emb)
  carry_gate = tf.sigmoid(carry_gate + carry_bias)
  if layer_norm:
    hidden_emb = contrib_layers.layer_norm(hidden_emb)

  output_emb = carry_gate * project_input_emb + (1 - carry_gate) * hidden_emb
  return output_emb
A single highway layer as described here: https://arxiv.org/abs/1505.00387.

  This is used as a generalization of a single hidden layer in a standard
  feedforward network. The difference is that the hidden layer with the
  non-linear activation is interpolated with a linear projection, enabling a
  "highway" for gradients.

  Args:
    input_emb: tensor<float> [..., embedding_size]
    output_size: An integer indicating the size of the hidden layer.
    dropout_ratio: The probability of dropping out each unit in the activation.
        This can be None, and is only applied during training.
    mode: One of the keys from tf.estimator.ModeKeys.
    layer_norm: Boolean indicating whether we should apply layer normalization.
    carry_bias: Float indicating the initial bias for the carry gate. This is
        typically set to a positive number to encourage information flow through
        the "highway" initially. This is analogous to the forget bias typically
        found in LSTMs.

  Returns:
    output_emb: A Tensor with the same shape as `input_emb`, except for the last
        dimension which will have size `output_size instead.
  
language/common/layers/common_layers.py