code: def get_kernel_name(
      is_mqa: bool, save_residuals: bool, is_segmented: bool, phase: str
  ) -> str:
      """Returns a unique name for all SplashAttention kernel variants."""
  
      assert phase == "dq" or phase == "dkv" or phase == "fwd"
      # Saving residuals is supported only for the fwd phase.
      assert not save_residuals or phase == "fwd"
      residuals = ""
      if save_residuals:
          residuals = "_residuals"
      elif phase == "fwd":
          residuals = "_no_residuals"
      attention_type = "mqa" if is_mqa else "mha"
      segments = "_segmented" if is_segmented else ""
      return f"splash_{attention_type}_{phase}{segments}{residuals}"
  
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 1 for calling a global function with example usage: get_kernel_name(is_mqa, save_residuals, is_segmented, phase) and returns: str, attention_type, str, phase, segments, residuals