code: def _splash_attention_fwd(
      fwd_mask_info: mask_info_lib.MaskInfo,
      dq_mask_info: mask_info_lib.MaskInfo | None,
      dkv_mask_info: mask_info_lib.MaskInfo | None,
      q: jax.Array,
      k: jax.Array,
      v: jax.Array,
      segment_ids: SegmentIds | None,
      save_residuals: bool,
      mask_value: float,
      is_mqa: bool,
      block_sizes: BlockSizes,
      residual_checkpoint_name: str | None,
      mask_function: MaskFunctionType | None,
      attn_logits_soft_cap: float | None = None,
  ) -> tuple[tuple[jax.Array], SplashResidualsType,]:
      if save_residuals:
          raise NotImplementedError("Higher-order AD not supported")
  
      out, (logsumexp,) = _splash_attention_forward(
          fwd_mask_info,
          q,
          k,
          v,
          segment_ids,
          mask_value=mask_value,
          is_mqa=is_mqa,
          block_sizes=block_sizes,
          save_residuals=True,
          attn_logits_soft_cap=attn_logits_soft_cap,
          residual_checkpoint_name=residual_checkpoint_name,
          mask_function=mask_function,
      )
      return out, (
          q,
          k,
          v,
          segment_ids,
          out,
          logsumexp,
          dq_mask_info,
          dkv_mask_info,
      )
  
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 20 for calling a global function with example usage: _splash_attention_fwd(fwd_mask_info, dq_mask_info, dkv_mask_info, q, k, v, segment_ids, save_residuals, mask_value, is_mqa, block_sizes, residual_checkpoint_name, mask_function, attn_logits_soft_cap) and returns: out,  q,  k,  v,  segment_ids,  out,  logsumexp,  dq_mask_info,  dkv_mask_info