code: def _flash_attention_dq_kernel(
      # Prefetched inputs
      data_next_ref,
      block_mask_ref,
      mask_next_ref,
      # Inputs
      q_ref,
      k_ref,
      v_ref,
      q_segment_ids_ref,
      kv_segment_ids_ref,
      logsumexp_ref,
      do_ref,
      di_ref,
      mask_ref,
      q_sequence_ref,
      # Outputs
      dq_scratch_ref,
      dq_ref,
      *,
      mask_value: float,
      grid_width: int,
      bq: int,
      bkv: int,
      attn_logits_soft_cap: float | None = None,
      q_layout: QKVLayout,
      k_layout: QKVLayout,
      v_layout: QKVLayout,
      mask_function: MaskFunctionType | None,
  ):
      float32 = jnp.float32
      HEAD_DIM_MINOR = QKVLayout.HEAD_DIM_MINOR
  
      h, i, j = pl.program_id(0), pl.program_id(1), pl.program_id(2)
  
      @pl.when(j == 0)
      def init():
          dq_scratch_ref[...] = jnp.zeros_like(dq_scratch_ref)
  
      _, _, should_run, should_not_mask = _next_nonzero(
          h, i, j, data_next_ref, block_mask_ref, mask_next_ref
      )
  
      @pl.when(should_run)
      def run():
          q = q_ref[...] if q_layout == HEAD_DIM_MINOR else q_ref[...].T
          # We keep k and v possibly transposed, since they are RHS of dots.
          k = k_ref[...]
          v = v_ref[...]
          logsumexp = jnp.expand_dims(logsumexp_ref[0], -1)
          do = do_ref[...]
          di = jnp.expand_dims(di_ref[0], -1)
  
          qk_dims = (
              NT_DIM_NUMBERS
              if k_layout == HEAD_DIM_MINOR
              else NN_DIM_NUMBERS
          )
          qk_uncapped = lax.dot_general(
              q, k, qk_dims, preferred_element_type=float32
          )
  
          qk = _apply_mask_and_soft_cap(
              qk_uncapped,
              mask_value,
              should_not_mask,
              mask_ref,
              q_sequence_ref,
              q_segment_ids_ref,
              kv_segment_ids_ref,
              attn_logits_soft_cap=attn_logits_soft_cap,
              k_slice=pl.ds(0, bkv),
              k_offset=j * bkv,
              bq=bq,
              mask_function=mask_function,
          )
          p = jnp.exp(qk - logsumexp)
          dp_dims = (
              NT_DIM_NUMBERS
              if v_layout == HEAD_DIM_MINOR
              else NN_DIM_NUMBERS
          )
          dp = lax.dot_general(
              do.astype(v.dtype),
              v,
              dp_dims,
              preferred_element_type=jnp.float32,
          )
          ds = (dp - di) * p
          if attn_logits_soft_cap is not None:
              normalized = qk_uncapped / attn_logits_soft_cap
              d = jnp.tanh(normalized)
              g = ds * (1 - d)
              ds = g + g * d
  
          dq_dims = (
              NN_DIM_NUMBERS
              if k_layout == HEAD_DIM_MINOR
              else NT_DIM_NUMBERS
          )
          dq_scratch_ref[...] += lax.dot_general(
              ds.astype(k.dtype),
              k,
              dq_dims,
              preferred_element_type=jnp.float32,
          )
  
      @pl.when(j == grid_width - 1)
      def end():
          dq_ref[...] = dq_scratch_ref[...].astype(dq_ref.dtype)
          dq_scratch_ref[...] = jnp.zeros_like(dq_scratch_ref)
  
imports: from __future__ import annotations
  import dataclasses
  import enum
  import functools
  from typing import Any, Callable, Literal, NamedTuple, Union, Optional, overload
  import jax
  from jax import ad_checkpoint
  from jax import lax
  from jax import tree_util
  from jax.experimental import pallas
  from jax.experimental.pallas import tpu
  from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask
  from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask_info
  import jax.numpy
  import numpy
label_desc: Use image node_id 21 for calling a global function with example usage: _flash_attention_dq_kernel(data_next_ref, block_mask_ref, mask_next_ref, q_ref, k_ref, v_ref, q_segment_ids_ref, kv_segment_ids_ref, logsumexp_ref, do_ref, di_ref, mask_ref, q_sequence_ref, dq_scratch_ref, dq_ref) without return types