code: def layer_norm_backward_kernel_dw_db(
      # Inputs
      x_ref,
      weight_ref,
      bias_ref,
      do_ref,
      mean_ref,
      rstd_ref,
      # Outputs
      dw_ref,
      db_ref,
      *,
      eps: float,
      block_m: int,
      block_n: int,
  ):
      m, n_col = x_ref.shape
      j = pl.program_id(0)
      col_idx = j * block_n + jnp.arange(block_n)
      col_mask = col_idx < n_col
  
      def body(i, acc_ref):
          row_idx = i * block_m + jnp.arange(block_m)
          row_mask = row_idx < m
          mask = row_mask[:, None] & col_mask[None, :]
          a = pl.load(
              x_ref,
              (row_idx[:, None], col_idx[None]),
              mask=mask,
              other=0.0,
          ).astype(jnp.float32)
          dout = pl.load(
              do_ref,
              (row_idx[:, None], col_idx[None]),
              mask=mask,
              other=0.0,
          ).astype(jnp.float32)
          mean = pl.load(
              mean_ref, (row_idx,), mask=row_mask, other=0.0
          ).astype(jnp.float32)
          rstd = pl.load(
              rstd_ref, (row_idx,), mask=row_mask, other=0.0
          ).astype(jnp.float32)
          a_hat = (a - mean[:, None]) * rstd[:, None]
          dw_acc_ref, db_acc_ref = acc_ref
          dw_acc_ref[:] += (dout * a_hat).sum(axis=0)
          db_acc_ref[:] += dout.sum(axis=0)
  
      dw_acc, db_acc = for_loop(
          pl.cdiv(m, block_m),
          body,
          (jnp.zeros(block_n), jnp.zeros(block_n)),
      )
      pl.store(
          dw_ref, (col_idx,), dw_acc.astype(dw_ref.dtype), mask=col_mask
      )
      pl.store(
          db_ref, (col_idx,), db_acc.astype(db_ref.dtype), mask=col_mask
      )
  
file_path: jax/jax/experimental/pallas/ops/layer_norm.py
imports: from __future__ import annotations
  import functools
  from typing import Optional
  import jax
  from jax import lax
  import jax.numpy
  from jax._src.lax.control_flow.for_loop import for_loop
  from jax.experimental import pallas
label_desc: Use image node_id 4 for calling a global function with example usage: layer_norm_backward_kernel_dw_db(x_ref, weight_ref, bias_ref, do_ref, mean_ref, rstd_ref, dw_ref, db_ref) without return types