code: def layer_norm(x, dim, epsilon=1e-6, name="layer_prepostprocess"):
    
    with tf.variable_scope(name + "/layer_norm"):
      scale = mtf.get_variable(
          x.mesh,
          "layer_norm_scale",
          mtf.Shape([dim]),
          initializer=tf.ones_initializer(),
          activation_dtype=x.dtype)
      bias = mtf.get_variable(
          x.mesh,
          "layer_norm_bias",
          mtf.Shape([dim]),
          initializer=tf.zeros_initializer(),
          activation_dtype=x.dtype)
      reduced_shape = x.shape - dim
      mean = mtf.reduce_mean(x, output_shape=reduced_shape)
      variance = mtf.reduce_mean(mtf.square(x - mean), output_shape=reduced_shape)
      norm_x = (x - mean) * mtf.rsqrt(variance + epsilon)
      return norm_x * scale + bias
docs: Layer normalization over dimension dim.
  
    Args:
      x: a mtf.Tensor whose shape contains dim.
      dim: a mtf.Dimension
      epsilon: a floating point number
      name: a string. variable scope.
  
    Returns:
      a mtf.Tensor with same shape as x.