code: def _flash_attention_bwd_dq(
      q,
      k,
      v,
      ab,
      segment_ids,
      l,
      m,
      do,
      di,
      *,
      block_q_major: int | None,
      block_k_major: int | None,
      block_k: int | None,
      sm_scale: float,
      causal: bool,
      mask_value: float,
      debug: bool,
  ):
      batch_size, num_heads, q_seq_len, head_dim = q.shape
      _, _, kv_seq_len, _ = k.shape
      _verify_block("block_q_dq", "q_seq_len", block_q_major, q_seq_len)
      _verify_block(
          "block_k_major_dq", "kv_seq_len", block_k_major, kv_seq_len
      )
      _verify_block("block_k_dq", "block_k", block_k, kv_seq_len)
  
      # Broadcast out scalar values
      m = jnp.broadcast_to(m[..., None], (*m.shape, MIN_BLOCK_SIZE))
      l = jnp.broadcast_to(l[..., None], (*l.shape, MIN_BLOCK_SIZE))
      # Preprocess contraction for bwd pass
      di = jnp.broadcast_to(di[..., None], (*di.shape, block_k_major))
  
      grid = (
          batch_size,
          num_heads,
          q_seq_len // block_q_major,
          kv_seq_len // block_k_major,
      )
  
      def qo_index_map(batch_index, head_index, q_seq_index, _):
          return (batch_index, head_index, q_seq_index, 0)
  
      qo_spec = pl.BlockSpec(
          qo_index_map, (1, 1, block_q_major, head_dim)
      )
      do_spec = qo_spec
  
      def kv_index_map(
          batch_index, head_index, q_seq_index, kv_seq_index
      ):
          if causal:
              # If the kv block is skipped, prefetch the next valid kv block, i.e. the
              # 0th one to be used for the next block_q rows.
              next_kv_index = lax.select(
                  below_or_on_diag(
                      q_seq_index,
                      block_q_major,
                      kv_seq_index,
                      block_k_major,
                  ),
                  kv_seq_index,
                  0,
              )
          else:
              next_kv_index = kv_seq_index
          return (batch_index, head_index, next_kv_index, 0)
  
      kv_spec = pl.BlockSpec(
          kv_index_map, (1, 1, block_k_major, head_dim)
      )
      assert kv_spec.block_shape is not None
      assert k.ndim == len(kv_spec.block_shape)
      assert v.ndim == len(kv_spec.block_shape)
  
      def lm_index_map(batch_index, head_index, q_seq_index, _):
          return (batch_index, head_index, q_seq_index, 0)
  
      lm_spec = pl.BlockSpec(
          lm_index_map, (1, 1, block_q_major, MIN_BLOCK_SIZE)
      )
      assert lm_spec.block_shape is not None
      assert l.ndim == len(lm_spec.block_shape)
      assert m.ndim == len(lm_spec.block_shape)
  
      di_spec = pl.BlockSpec(
          qo_index_map, (1, 1, block_q_major, MIN_BLOCK_SIZE)
      )
      assert di_spec.block_shape is not None
      assert di.ndim == len(di_spec.block_shape)
  
      def ab_index_map(
          batch_index, head_index, q_seq_index, kv_seq_index
      ):
          return (batch_index, head_index, q_seq_index, kv_seq_index)
  
      dab_spec = (
          pl.BlockSpec(
              ab_index_map, (1, 1, block_q_major, block_k_major)
          )
          if ab is not None
          else None
      )
  
      q_segment_ids_spec = kv_segment_ids_spec = None
      q_segment_ids = kv_segment_ids = None
      if segment_ids is not None:
  
          def q_segment_ids_index_map(
              batch_index, head_index, q_seq_index, _
          ):
              del head_index
              return (batch_index, q_seq_index, 0)
  
          def kv_segment_ids_index_map(
              batch_index, head_index, q_seq_index, kv_seq_index
          ):
              del head_index
              if causal:
                  # If the kv block is skipped, prefetch the next valid kv block, i.e. the
                  # 0th one to be used for the next block_q rows.
                  next_kv_index = lax.select(
                      below_or_on_diag(
                          q_seq_index,
                          block_q_major,
                          kv_seq_index,
                          block_k_major,
                      ),
                      kv_seq_index,
                      0,
                  )
              else:
                  next_kv_index = kv_seq_index
              return (batch_index, 0, next_kv_index)
  
          q_segment_ids_spec = pl.BlockSpec(
              q_segment_ids_index_map, (1, block_q_major, NUM_LANES)
          )
          kv_segment_ids_spec = pl.BlockSpec(
              kv_segment_ids_index_map, (1, NUM_SUBLANES, block_k_major)
          )
  
          q_segment_ids = jax.lax.broadcast_in_dim(
              segment_ids.q,
              (batch_size, q_seq_len, NUM_LANES),
              (
                  0,
                  1,
              ),
          )
          kv_segment_ids = jax.lax.broadcast_in_dim(
              segment_ids.kv,
              (batch_size, NUM_SUBLANES, kv_seq_len),
              (
                  0,
                  2,
              ),
          )
  
      in_specs = [
          qo_spec,
          kv_spec,
          kv_spec,
          dab_spec,
          q_segment_ids_spec,
          kv_segment_ids_spec,
          lm_spec,
          lm_spec,
          do_spec,
          di_spec,
      ]
  
      out_shapes = [
          jax.ShapeDtypeStruct(q.shape, q.dtype),
          jax.ShapeDtypeStruct((block_q_major, head_dim), jnp.float32),
          jax.ShapeDtypeStruct(ab.shape, ab.dtype)
          if ab is not None
          else None,
      ]
      dq_spec = pl.BlockSpec(
          qo_index_map, (1, 1, block_q_major, head_dim)
      )
      out_specs = [
          dq_spec,
          pl.BlockSpec(lambda *_: (0, 0), (block_q_major, head_dim)),
          dab_spec,
      ]
  
      kernel = functools.partial(
          _flash_attention_dq_kernel,
          sm_scale=sm_scale,
          causal=causal,
          mask_value=mask_value,
          block_k=block_k,
          kv_seq_len=kv_seq_len,
      )
      name_scope = f"flash_mha_bwd_dq_{block_q_major=}_{block_k_major=}_{block_k=}"
      with jax.named_scope(name_scope):
          dq, _, ds = pl.pallas_call(
              kernel,
              in_specs=in_specs,  # type: ignore
              out_shape=out_shapes,
              out_specs=out_specs,  # type: ignore
              grid=grid,
              debug=debug,
              mosaic_params=dict(
                  dimension_semantics=(
                      "parallel",
                      "parallel",
                      "parallel",
                      "arbitrary",
                  )
              ),
          )(q, k, v, ab, q_segment_ids, kv_segment_ids, l, m, do, di)
  
      # dab is just ds
      return dq, ds
  
imports: from __future__ import annotations
  import dataclasses
  import functools
  from typing import Any, NamedTuple
  import jax
  from jax import lax
  from jax.experimental import pallas
  from jax.experimental.pallas import tpu
  import jax.numpy
label_desc: Use image node_id 13 for calling a global function with example usage: _flash_attention_bwd_dq(q, k, v, ab, segment_ids, l, m, do, di) and returns: dq,  ds, batch_index,  head_index,  q_seq_index,  int, batch_index,  head_index,  next_kv_index,  int, batch_index,  head_index,  q_seq_index,  int, batch_index,  head_index,  q_seq_index,  kv_seq_index, batch_index,  q_seq_index,  int, batch_index,  int,  next_kv_index