code: def norm_forward(
      x: Tensor,
      mean: Tensor,
      sqr_mean: Tensor,
      weight: Tensor,
      bias: Tensor,
      eps: float,
  ):
      mu = x - mean
      var = sqr_mean - mean**2
      sigma = torch.sqrt(var + eps)
      z = mu / sigma
      output = weight * z + bias
  
      return output, mu, sigma
  
imports: from typing import Optional, Tuple
  import torch
  from torch import Tensor
  from torch.cuda.amp import custom_bwd, custom_fwd
  from colossalai.legacy.communication import all_gather, all_reduce, broadcast, reduce, reduce_scatter
  from colossalai.legacy.constants import INPUT_GROUP_3D, WEIGHT_GROUP_3D
  from colossalai.legacy.context.parallel_mode import ParallelMode
  from colossalai.legacy.core import global_context
  from ._utils import get_parallel_mode_from_env, push_async_grad
label_desc: Use image node_id 4 for calling a global function with example usage: norm_forward(x, mean, sqr_mean, weight, bias, eps) and returns: output,  mu,  sigma