code: def _attention_reference_custom_fwd(
      mask: jax.Array,  # [q_seq_len, kv_seq_len]
      q: jax.Array,  # [q_seq_len, head_dim]
      k: jax.Array,  # [kv_seq_len, head_dim]
      v: jax.Array,  # [kv_seq_len, head_dim]
      segment_ids: SegmentIds | None,
      mask_value: float,
      save_residuals: bool,
      custom_type: str,
      attn_logits_soft_cap: float | None,
  ):
      if save_residuals:
          raise NotImplementedError("Higher-order AD not supported.")
  
      o, (logsumexp,) = _attention_reference(
          mask,
          q,
          k,
          v,
          segment_ids,
          mask_value=mask_value,
          save_residuals=True,
          custom_type=custom_type,
          attn_logits_soft_cap=attn_logits_soft_cap,
      )
      return o, (mask, q, k, v, segment_ids, o, logsumexp)
  
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 7 for calling a global function with example usage: _attention_reference_custom_fwd(mask, q, k, v, segment_ids, mask_value, save_residuals, custom_type, attn_logits_soft_cap) and returns: o,  mask,  q,  k,  v,  segment_ids,  o,  logsumexp