code: def _preprocess_backward_kernel(
      out_ref, dout_ref, l_ref, new_dout_ref, delta_ref, *, block_q: int
  ):
      pid_m = pl.program_id(0)
  
      off_m = pl.ds(pid_m * block_q, block_q)
      # load
      o = pl.load(out_ref, (off_m, slice(None))).astype(jnp.float32)
      do = pl.load(dout_ref, (off_m, slice(None))).astype(jnp.float32)
      denom = pl.load(l_ref, (off_m,)).astype(jnp.float32)
      # compute
      do = do / denom[:, None]
      delta = jnp.sum(o * do, axis=1)
      # write-back
      pl.store(
          new_dout_ref,
          (off_m, slice(None)),
          do.astype(new_dout_ref.dtype),
      )
      pl.store(delta_ref, (off_m,), delta.astype(delta_ref.dtype))
  
file_path: jax/jax/experimental/pallas/ops/attention.py
imports: from __future__ import annotations
  import functools
  from typing import Any, Optional
  import jax
  from jax import lax
  from jax.experimental import pallas
  import jax.numpy
  import numpy
label_desc: Use image node_id 5 for calling a global function with example usage: _preprocess_backward_kernel(out_ref, dout_ref, l_ref, new_dout_ref, delta_ref) without return types