Connects the LayerNorm module into the graph.

    Args:
      inputs: a Tensor of dimensionality >= 2.

    Returns:
      normalized: layer normalized outputs with same shape as inputs.

    Raises:
      base.NotSupportedError: If `inputs` has less than 2 dimensions.

def _build(self, inputs):
    """Connects the LayerNorm module into the graph.

    Args:
      inputs: a Tensor of dimensionality >= 2.

    Returns:
      normalized: layer normalized outputs with same shape as inputs.

    Raises:
      base.NotSupportedError: If `inputs` has less than 2 dimensions.
    """

    if self._axis is None:
      axis = list(range(1, inputs.shape.ndims))
    else:
      axis = self._axis

    original_dtype = inputs.dtype
    if original_dtype in [tf.float16, tf.bfloat16]:
      inputs = tf.cast(inputs, tf.float32)

    if inputs.get_shape().ndims < 2:
      raise base.NotSupportedError(
          "Layer normalization expects inputs of at least rank 2."
          " Got inputs of rank {}.".format(inputs.get_shape().ndims))

    # Shape for the learnable scale and offset is the number of channels. See
    # https://arxiv.org/pdf/1803.08494.pdf around equation 6.
    params_shape = inputs.get_shape()[-1:]

    if self._scale:
      if self.GAMMA not in self._initializers:
        self._initializers[self.GAMMA] = create_gamma_initializer()
      self._gamma = tf.get_variable(
          self.GAMMA,
          shape=params_shape,
          dtype=inputs.dtype,
          initializer=self._initializers[self.GAMMA],
          partitioner=self._partitioners.get(self.GAMMA),
          regularizer=self._regularizers.get(self.GAMMA))
    else:
      self._gamma = None

    if self._offset:
      if self.BETA not in self._initializers:
        self._initializers[self.BETA] = create_beta_initializer()
      self._beta = tf.get_variable(
          self.BETA,
          shape=params_shape,
          dtype=inputs.dtype,
          initializer=self._initializers[self.BETA],
          partitioner=self._partitioners.get(self.BETA),
          regularizer=self._regularizers.get(self.BETA))
    else:
      self._beta = None

    mean, var = tf.nn.moments(inputs, axis, keep_dims=True)

    normalized = tf.nn.batch_normalization(inputs, mean, var, self._beta,
                                           self._gamma, self._eps)

    if original_dtype in [tf.float16, tf.bfloat16]:
      normalized = tf.cast(normalized, dtype=original_dtype)
    return normalized