def layer_norm(layer_inputs, hidden_size):
  """Implements layer norm from [Ba et al. 2016] Layer Normalization.

  See eqn. 4 in (https://arxiv.org/pdf/1607.06450.pdf).

  Args:
    layer_inputs (tensor): The inputs to the layer.
      shape <float32>[batch_size, hidden_size]
    hidden_size (int): Dimensionality of the hidden layer.

  Returns:
    normalized (tensor): layer_inputs, normalized over all the hidden units in
      the layer.
      shape <float32>[batch_size, hidden_size]
  """

  mean, var = tf.nn.moments(layer_inputs, [1], keep_dims=True)
  with tf.variable_scope("layernorm", reuse=tf.AUTO_REUSE):
    gain = tf.get_variable(
        "gain", shape=[hidden_size], initializer=tf.constant_initializer(1))
    bias = tf.get_variable(
        "bias", shape=[hidden_size], initializer=tf.constant_initializer(0))

  normalized = gain * (layer_inputs - mean) / tf.sqrt(var + _EPSILON) + bias
  return normalized
Implements layer norm from [Ba et al. 2016] Layer Normalization.

  See eqn. 4 in (https://arxiv.org/pdf/1607.06450.pdf).

  Args:
    layer_inputs (tensor): The inputs to the layer.
      shape <float32>[batch_size, hidden_size]
    hidden_size (int): Dimensionality of the hidden layer.

  Returns:
    normalized (tensor): layer_inputs, normalized over all the hidden units in
      the layer.
      shape <float32>[batch_size, hidden_size]
  
language/labs/memory/model_utils.py