code: def _layer_norm_backward_dw_db_partial_strided(
      dy_ptr,
      x_ptr,
      x_invstd_ptr,
      x_mean_ptr,
      dw_partial_buf_ptr,
      db_partial_buf_ptr,
      M: tl.constexpr,
      N: tl.constexpr,
      BUF_N_STRIDE: tl.constexpr,
      N_BLOCK: tl.constexpr,
      M_PARTIAL_REDUCE: tl.constexpr,
      D0: tl.constexpr,
      D1: tl.constexpr,
      D2: tl.constexpr,
      D3: tl.constexpr,
      S0: tl.constexpr,
      S1: tl.constexpr,
      S2: tl.constexpr,
      S3: tl.constexpr,
  ):
      m_logic_idx = tl.program_id(0) * M_PARTIAL_REDUCE + tl.arange(
          0, M_PARTIAL_REDUCE
      )
      m_mask = m_logic_idx < M
      m_logic_idx_0 = m_logic_idx // (D1 * D2) % D0
      m_logic_idx_1 = m_logic_idx // D2 % D1
      m_logic_idx_2 = m_logic_idx % D2
      m_idx = (
          m_logic_idx_0 * S0 + m_logic_idx_1 * S1 + m_logic_idx_2 * S2
      )
      n_logic_idx = tl.program_id(1) * N_BLOCK + tl.arange(0, N_BLOCK)
      n_mask = n_logic_idx < N
      n_idx = n_logic_idx * S3
      mask = m_mask[:, None] & n_mask[None, :]
      x_idx = m_idx[:, None] + n_idx[None, :]
      x = tl.load(x_ptr + x_idx, mask, other=0).to(tl.float32)
      x_mean = tl.load(x_mean_ptr + m_logic_idx, m_mask, other=0).to(
          tl.float32
      )[:, None]
      x_invstd = tl.load(
          x_invstd_ptr + m_logic_idx, m_mask, other=0
      ).to(tl.float32)[:, None]
      x_hat = (x - x_mean) * x_invstd
      dy_idx = N * m_logic_idx[:, None] + n_logic_idx[None, :]
      dy = tl.load(dy_ptr + dy_idx, mask, other=0).to(tl.float32)
      dw_partial = tl.sum(dy * x_hat, axis=0)
      db_partial = tl.sum(dy, axis=0)
      tl.store(
          dw_partial_buf_ptr
          + BUF_N_STRIDE * n_logic_idx
          + tl.program_id(0),
          dw_partial,
          n_mask,
      )
      tl.store(
          db_partial_buf_ptr
          + BUF_N_STRIDE * n_logic_idx
          + tl.program_id(0),
          db_partial,
          n_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 4 for calling a global function with example usage: _layer_norm_backward_dw_db_partial_strided(dy_ptr, x_ptr, x_invstd_ptr, x_mean_ptr, dw_partial_buf_ptr, db_partial_buf_ptr, M, N, BUF_N_STRIDE, N_BLOCK, M_PARTIAL_REDUCE, D0, D1, D2, D3, S0, S1, S2, S3) without return types