code: def _attention_reference_default(
      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,
  ):
      del custom_type
      logits = jnp.einsum(
          "sd,td->st", q.astype(jnp.float32), k.astype(jnp.float32)
      )
  
      if segment_ids is not None:
          mask = jnp.logical_and(
              mask, segment_ids.q[:, None] == segment_ids.kv[None, :]
          )
  
      if attn_logits_soft_cap is not None:
          logits = jnp.tanh(logits / attn_logits_soft_cap)
          logits = logits * attn_logits_soft_cap
  
      logits = jnp.where(mask, logits, mask_value)
      m = logits.max(axis=-1)
      s = jnp.exp(logits - m[..., None])
      l = s.sum(axis=-1)
      s = s / l[..., None]
  
      o = jnp.einsum("st,td->sd", s, v.astype(jnp.float32))
  
      logsumexp = m + jnp.log(l)
      if save_residuals:
          return o, (logsumexp,)
      return o
  
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 5 for calling a global function with example usage: _attention_reference_default(mask, q, k, v, segment_ids, mask_value, save_residuals, custom_type, attn_logits_soft_cap) and returns: o, o,  logsumexp