code: def layer_norm(a, weight, bias, eps):
      assert a.is_contiguous()
      assert weight.is_contiguous()
      assert bias.is_contiguous()
  
      # allocate output
      out = torch.empty_like(a)
      # reshape input data into 2D tensor
      a_arg = a.view(-1, a.shape[-1])
      M, N = a_arg.shape
      # Less than 64KB per feature: enqueue fused kernel
      MAX_FUSED_SIZE = 65536 // a.element_size()
      BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(N))
      BLOCK_SIZE = max(BLOCK_SIZE, 128)
      BLOCK_SIZE = min(BLOCK_SIZE, 4096)
      BLOCK_SIZE = BLOCK_SIZE if N <= 4096 else 8192
      # heuristics for number of warps
      num_warps = min(max(BLOCK_SIZE // 256, 1), 8)
      layer_norm_kernel[(M,)](
          out,
          a_arg,
          weight,
          bias,
          a_arg.stride(0),
          N,
          eps,
          BLOCK_SIZE=BLOCK_SIZE,
          num_warps=num_warps,
      )
      return out
  
file_path: deepspeed/deepspeed/ops/transformer/inference/triton/layer_norm.py
imports: import torch
  import triton
  import triton.language
label_desc: Use image node_id 4 for calling a global function with example usage: layer_norm(a, weight, bias, eps) and returns: out