code: def _attention_reference_custom_bwd(
      mask_value: float,
      save_residuals: bool,
      custom_type: str,
      attn_logits_soft_cap: float | None,
      res,
      do: jax.Array,
  ) -> tuple[None, jax.Array, jax.Array, jax.Array, None]:
      del save_residuals
      mask, q, k, v, segment_ids, o, logsumexp = res
  
      uncapped_logits = jnp.einsum(
          "qc,kc->qk", q, k, preferred_element_type=jnp.float32
      )
  
      if attn_logits_soft_cap is not None:
          logits = jnp.tanh(uncapped_logits / attn_logits_soft_cap)
          logits = logits * attn_logits_soft_cap
      else:
          logits = uncapped_logits
  
      if segment_ids is not None:
          mask = jnp.logical_and(
              mask, segment_ids.q[:, None] == segment_ids.kv[None, :]
          )
      logits = jnp.where(mask, logits, mask_value)
  
      p = jnp.exp(logits - logsumexp[..., None])
      do = do.astype(jnp.float32)  # pytype: disable=attribute-error
      dv = jnp.einsum("pt,pd->td", p, do).astype(v.dtype)
      dp = jnp.einsum("pd,td->pt", do, v.astype(jnp.float32))
  
      # These two ways of computing ds are mathematically equivalent. The first
      # involves reducing over the head_dim dimension and the second involves
      # reducing over a sequence dimension. They tend to produce slightly different
      # numerics.
      if custom_type == "flash":
          di = jnp.sum(o.astype(jnp.float32) * do, axis=-1)[..., None]
      else:
          di = jnp.einsum("st,st->s", dp, p)[:, None]
      ds = (dp - di) * p
      if attn_logits_soft_cap is not None:
          normalized = uncapped_logits / attn_logits_soft_cap
          d = jnp.tanh(normalized)
          g = ds * (1 - d)
          ds = g + g * d
      dk = jnp.einsum("sd,st->td", q.astype(jnp.float32), ds).astype(
          k.dtype
      )
      dq = jnp.einsum("st,td->sd", ds, k.astype(jnp.float32)).astype(
          q.dtype
      )
      return None, dq, dk, dv, None
  
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 8 for calling a global function with example usage: _attention_reference_custom_bwd(mask_value, save_residuals, custom_type, attn_logits_soft_cap, res, do) and returns: None,  dq,  dk,  dv,  None