code: def test_splash_attention_fwd(self, is_mqa, is_segmented, data):
      seed = data.draw(seed_strategy())
      key = random.key(seed)
      k1, k2, k3 = random.split(key, 3)
  
      (
          q_seq_len,
          kv_seq_len,
          num_q_heads,
          num_kv_heads,
          head_dim,
          dtype,
      ) = data.draw(mha_strategy())
  
      # Avoid segment ids for rectangular matrices, as its hard to enforce
      # valid masks (non-0 rows).
      hp.assume(q_seq_len == kv_seq_len or not is_segmented)
  
      q = random.uniform(
          k1, (num_q_heads, q_seq_len, head_dim), dtype=dtype
      )
      if is_mqa:
          k = random.uniform(k2, (kv_seq_len, head_dim), dtype=dtype)
          v = random.uniform(k3, (kv_seq_len, head_dim), dtype=dtype)
      else:
          k = random.uniform(
              k2, (num_kv_heads, kv_seq_len, head_dim), dtype=dtype
          )
          v = random.uniform(
              k3, (num_kv_heads, kv_seq_len, head_dim), dtype=dtype
          )
  
      segment_ids = None
      if is_segmented:
          assert q_seq_len == kv_seq_len
          segment_ids = data.draw(segment_ids_strategy(q_seq_len))
      attn_logits_soft_cap = data.draw(attn_logits_soft_cap_strategy())
      masks = data.draw(
          mha_mask_strategy(q_seq_len, kv_seq_len, num_q_heads)
      )
      mask = mask_lib.MultiHeadMask(
          tuple((m.get_mask() for m in masks))
      )
      block_sizes = data.draw(
          block_sizes_strategy(q_seq_len, kv_seq_len)
      )
      if is_mqa:
          attn_ref = splash.make_masked_mqa_reference(mask)
          attn = splash.make_splash_mqa_single_device(
              mask,
              block_sizes=block_sizes,
              save_residuals=True,
              attn_logits_soft_cap=attn_logits_soft_cap,
          )
      else:
          attn_ref = splash.make_masked_mha_reference(mask)
          attn = splash.make_splash_mha_single_device(
              mask,
              block_sizes=block_sizes,
              save_residuals=True,
              attn_logits_soft_cap=attn_logits_soft_cap,
          )
      attn_ref = partial(
          attn_ref,
          save_residuals=True,
          attn_logits_soft_cap=attn_logits_soft_cap,
      )
      o, (logsumexp,) = attn(q, k, v, segment_ids)
      o_ref, (logsumexp_ref,) = attn_ref(
          q.astype(jnp.float32),
          k.astype(jnp.float32),
          v.astype(jnp.float32),
          segment_ids,
      )
      self._assert_allclose(o, o_ref, atol=3e-3, rtol=3e-3)
      self._assert_allclose(
          logsumexp, logsumexp_ref, atol=1e-3, rtol=1e-3
      )
  
imports: from __future__ import annotations
  import dataclasses
  import functools
  from typing import Any, Callable, TypeVar
  from absl.testing import absltest
  from absl.testing import parameterized
  import hypothesis
  import hypothesis.strategies
  import jax
  from jax import random
  from jax import tree_util
  from jax._src import test_util
  from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_kernel
  from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask
  import jax.numpy
  import numpy
label_desc: Use image node_id 2 for calling the SplashAttentionTest obj's underlying member method code with example usage: obj.test_splash_attention_fwd(is_mqa, is_segmented, data) without return types