code: def __call__(self, *args, **kwargs) -> SplashCustomReturnType:
      return _splash_attention(
          self.fwd_mask_info,
          self.dq_mask_info,
          self.dkv_mask_info,
          *args,
          **kwargs,
          **self.kwargs,
      )
  
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 2 for calling the SplashAttentionKernel obj's underlying member method code with example usage: obj.__call__() and returns: _splash_attention