code: def _mha_reference_bwd(
      causal: bool,
      mask_value: float,
      sm_scale: float,
      save_residuals: bool,
      residuals,
      do,
  ):
      del save_residuals
      q, k, v, ab, segment_ids, o, l, m = residuals
      dq, dk, dv, dab = mha_reference_bwd(
          q,
          k,
          v,
          ab,
          segment_ids,
          o,
          l,
          m,
          do,
          causal=causal,
          mask_value=mask_value,
          sm_scale=sm_scale,
      )
      return dq, dk, dv, dab, 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 19 for calling a global function with example usage: _mha_reference_bwd(causal, mask_value, sm_scale, save_residuals, residuals, do) and returns: dq,  dk,  dv,  dab,  None