code: def _splash_attention_custom(
      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,
  ) -> SplashCustomReturnType:
      # The forward function does not use the dq and dkv MaskInfos, it just forwards
      # them to the backward function as residuals. This is a way to communicate
      # arbitrary Arrays to the backward function. Since the three MaskInfos are
      # constants there is no overhead in passing them to the backward function as
      # residuals. When sharding computation MaskInfos are partitioned so both the
      # forward and the backward kernels need to work on the relevant slice. If we
      # recomputed the backward MaskInfos in the backward function from the numpy
      # mask then we would not work with the MaskInfo slice relevant to the current
      # device.
      del dq_mask_info, dkv_mask_info
  
      return (
          _splash_attention_forward(  # pytype: disable=wrong-arg-types
              fwd_mask_info,
              q,
              k,
              v,
              segment_ids,
              mask_value=mask_value,
              is_mqa=is_mqa,
              block_sizes=block_sizes,
              save_residuals=save_residuals,
              attn_logits_soft_cap=attn_logits_soft_cap,
              residual_checkpoint_name=residual_checkpoint_name,
              mask_function=mask_function,
          )
      )
  
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 19 for calling a global function with example usage: _splash_attention_custom(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: _splash_attention_forward