code: def _flash_attention_bwd(
      save_residuals: bool,
      causal: bool,
      sm_scale: float,
      block_sizes: BlockSizes,
      debug: bool,
      residuals,
      do,
  ):
      """VJP rule for FlashAttention."""
      if save_residuals:
          raise NotImplementedError("Higher-order AD not supported")
      (q, k, v, ab, segment_ids, o, l, m) = residuals
      if not block_sizes.has_backward_blocks:
          raise ValueError(
              "Program is being differentiated, but not all backward blocks are"
              " specified"
          )
  
      di = jnp.sum(
          o.astype(jnp.float32) * do.astype(jnp.float32), axis=-1
      )  # [batch_size, num_heads, q_seq_len]
  
      dk, dv = _flash_attention_bwd_dkv(
          q,
          k,
          v,
          ab,
          segment_ids,
          l,
          m,
          do,
          di,
          block_q_major=block_sizes.block_q_major_dkv,
          block_k_major=block_sizes.block_k_major_dkv,
          block_k=block_sizes.block_k_dkv,
          block_q=block_sizes.block_q_dkv,
          sm_scale=sm_scale,
          causal=causal,
          mask_value=DEFAULT_MASK_VALUE,
          debug=debug,
      )
  
      dq, ds = _flash_attention_bwd_dq(
          q,
          k,
          v,
          ab,
          segment_ids,
          l,
          m,
          do,
          di,
          block_q_major=block_sizes.block_q_dq,
          block_k_major=block_sizes.block_k_major_dq,
          block_k=block_sizes.block_k_dq,
          sm_scale=sm_scale,
          causal=causal,
          mask_value=DEFAULT_MASK_VALUE,
          debug=debug,
      )
      return dq, dk, dv, ds, None
  
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 4 for calling a global function with example usage: _flash_attention_bwd(save_residuals, causal, sm_scale, block_sizes, debug, residuals, do) and returns: dq,  dk,  dv,  ds,  None