code: def layer_norm(
      x,
      dim,
      epsilon=1e-6,
      subtract_mean=True,
      use_scale=True,
      use_bias=True,
      name=None,
  ):
      """Layer normalization over dimension dim.
  
      TODO(noam): This is cleaner than the version in mtf.layers
      Move this version into mtf.layers to replace the one there.
  
      Args:
        x: a mtf.Tensor whose shape contains dim.
        dim: a mtf.Dimension
        epsilon: a floating point number
        subtract_mean: a boolean
        use_scale: a boolean
        use_bias: a boolean
        name: a string used for tf.variable_scope.
  
      Returns:
        a mtf.Tensor with same shape as x.
      """
      with tf.variable_scope(name, default_name="layer_norm"):
          if subtract_mean:
              x -= mtf.reduce_mean(x, reduced_dim=dim)
          variance = mtf.reduce_mean(mtf.square(x), reduced_dim=dim)
          x *= mtf.rsqrt(variance + epsilon)
          if use_scale:
              x *= mtf.get_variable(
                  x.mesh,
                  "scale",
                  mtf.Shape([dim]),
                  initializer=tf.ones_initializer(),
                  activation_dtype=x.dtype,
              )
          if use_bias:
              x += mtf.get_variable(
                  x.mesh,
                  "bias",
                  mtf.Shape([dim]),
                  initializer=tf.zeros_initializer(),
                  activation_dtype=x.dtype,
              )
          return x
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import copy
  import json
  import math
  import re
  import mesh_tensorflow
  import mesh_tensorflow.transformer.moe
  import six
  import tensorflow.compat.v1
label_desc: Use image node_id 4 for calling a global function with example usage: layer_norm(x, dim, epsilon, subtract_mean, use_scale, use_bias, name) and returns: x