code: def _test_impl(
      S, B, hidden_size, itype, wtype, ctype=fp32, mem_eff=False
  ):
      seed = 1243
      torch.manual_seed(seed)
      torch.cuda.manual_seed(seed)
  
      otype = wtype
      print("========================================================")
      print(
          f"S={S} B={B} Hidden={hidden_size} {itype} {wtype} Mem_Eff={mem_eff}"
      )
      print("--------------------------------------------------------")
  
      x = torch.randn(S * B, hidden_size, dtype=itype, device=device)
      gamma = torch.randn(hidden_size, dtype=wtype, device=device) * 0.2
      beta = torch.randn(hidden_size, dtype=wtype, device=device) * 0.2
      epsilon = 1e-5
  
      x.requires_grad = True
      gamma.requires_grad = True
      beta.requires_grad = True
  
      mu_ref = x.mean(1, dtype=ctype, keepdim=True)
      v = torch.square(x - mu_ref).mean(1, dtype=ctype, keepdim=True)
      rs_ref = torch.rsqrt(v + epsilon)
      y_ref = rs_ref * (x.to(ctype) - mu_ref)
      z_ref = (
          gamma.unsqueeze(0) * (y_ref).to(otype) + beta.unsqueeze(0)
      ).to(otype)
  
      mu_ref = mu_ref.flatten()
      rs_ref = rs_ref.flatten()
  
      dz = torch.randn_like(z_ref)
  
      # z_ref.backward(dz)
      # dx_ref = x.grad
      # dgamma_ref = gamma.grad
      # dbeta_ref = beta.grad
  
      dx_ref, dg_ref, db_ref = backward_(dz, x, mu_ref, rs_ref, gamma)
  
      z, mu, rs = fln.ln_fwd(x, gamma, beta, epsilon)
      if mem_eff:
          dx, dg, db, dg_part, db_part = fln.ln_bwd(
              dz, z, mu, rs, gamma, beta, True
          )
      else:
          dx, dg, db, dg_part, db_part = fln.ln_bwd(
              dz, x, mu, rs, gamma, beta, False
          )
  
      re_z, mse_z = metrics(z_ref, z)
      re_mu, mse_mu = metrics(mu_ref, mu)
      re_rs, mse_rs = metrics(rs_ref, rs)
  
      re_dx, mse_dx = metrics(dx_ref, dx)
      re_dg, mse_dg = metrics(dg_ref, dg)
      re_db, mse_db = metrics(db_ref, db)
  
      print(f" z: relerr={re_z :.4e} mse={mse_z :.4e}")
      print(f"mu: relerr={re_mu:.4e} mse={mse_mu:.4e}")
      print(f"rs: relerr={re_mu:.4e} mse={mse_mu:.4e}")
  
      print(f"dx: relerr={re_dx:.4e} mse={mse_dx:.4e}")
      print(f"dg: relerr={re_dg:.4e} mse={mse_dg:.4e}")
      print(f"db: relerr={re_db:.4e} mse={mse_db:.4e}")
  
      def check_err(x, relerr):
          tol = (
              2e-2
              if x.dtype in (torch.float16, torch.bfloat16)
              else 5e-6
          )
          return relerr < tol
  
      return [
          check_err(x, re)
          for x, re in zip(
              [z, mu, rs, dx, dg, db],
              [re_z, re_mu, re_rs, re_dx, re_dg, re_db],
          )
      ]
  
file_path: apex/apex/contrib/test/layer_norm/test_fast_layer_norm.py
imports: import itertools
  import unittest
  import torch
label_desc: Use image node_id 5 for calling a global function with example usage: _test_impl(S, B, hidden_size, itype, wtype, ctype, mem_eff) and returns: unknown, unknown