code: def _layer_norm_backward_dx(
      dy_ptr,
      x_ptr,
      w_ptr,
      x_invstd_ptr,
      x_mean_ptr,
      dx_ptr,
      M: tl.constexpr,
      N: tl.constexpr,
      M_BLOCK: tl.constexpr,
      N_BLOCK: tl.constexpr,
  ):
      m_idx = (tl.program_id(0) * M_BLOCK + tl.arange(0, M_BLOCK))[
          :, None
      ]
      m_mask = m_idx < M
      n_idx = tl.arange(0, N_BLOCK)[None, :]
      n_mask = n_idx < N
      mask = m_mask & n_mask
      x = tl.load(x_ptr + N * m_idx + n_idx, mask, other=0).to(
          tl.float32
      )
      x_mean = tl.load(x_mean_ptr + m_idx, m_mask, other=0).to(
          tl.float32
      )
      x_invstd = tl.load(x_invstd_ptr + m_idx, m_mask, other=0).to(
          tl.float32
      )
      x_hat = (x - x_mean) * x_invstd
      dy = tl.load(dy_ptr + N * m_idx + n_idx, mask, other=0).to(
          tl.float32
      )
      w = tl.load(w_ptr + n_idx, n_mask, other=0).to(tl.float32)
      c1 = tl.sum(x_hat * dy * w, axis=1) / N
      c2 = tl.sum(dy * w, axis=1) / N
      dx = x_invstd * (dy * w - c1[:, None] * x_hat - c2[:, None])
      tl.store(dx_ptr + N * m_idx + n_idx, dx, mask)
  
file_path: apex/apex/contrib/openfold_triton/_layer_norm_backward_kernels.py
imports: import torch
  import triton
  import triton.language
  from triton import Config
label_desc: Use image node_id 1 for calling a global function with example usage: _layer_norm_backward_dx(dy_ptr, x_ptr, w_ptr, x_invstd_ptr, x_mean_ptr, dx_ptr, M, N, M_BLOCK, N_BLOCK) without return types