docs: 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