docstr: Returns a value that can be used as a shard_map partition spec for the kernel.
code: def manual_sharding_spec(self, sharding: jax.sharding.NamedSharding):
      """Returns a value that can be used as a shard_map partition spec for the kernel."""
      if self.fwd_mask_info.data_next is not None:
          block_mask_shape = self.fwd_mask_info.data_next.shape
          try:
              shard_shape = sharding.shard_shape(block_mask_shape)
          except ValueError as exc:
              raise ValueError(
                  "The sharding must divide the mask blocks evenly between devices"
              ) from exc
          if block_mask_shape[-1] != shard_shape[-1]:
              raise ValueError(
                  "Sharding the kv sequence dimension is not supported"
              )
      spec = sharding.spec
      assert len(spec) == 2
      replicated = jax.sharding.PartitionSpec()
      # Shard q_sequence over the sequence dimension only.
      q_sequence_spec = jax.sharding.PartitionSpec(spec[1])
      mask_info_specs = (
          mask_info_lib.MaskInfo(  # pytype: disable=wrong-arg-types
              data_next=spec
              if self.fwd_mask_info.data_next is not None
              else None,
              mask_next=spec
              if self.fwd_mask_info.mask_next is not None
              else None,
              block_mask=spec
              if self.fwd_mask_info.block_mask is not None
              else None,
              partial_mask_blocks=replicated
              if self.fwd_mask_info.partial_mask_blocks is not None
              else None,
              q_sequence=q_sequence_spec
              if self.fwd_mask_info.q_sequence is not None
              else None,
          )
      )
      return SplashAttentionKernel(
          mask_info_specs,
          mask_info_specs if self.dq_mask_info is not None else None,
          mask_info_specs if self.dkv_mask_info is not None else None,
          **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 3 for calling the SplashAttentionKernel obj's underlying member method code with example usage: obj.manual_sharding_spec(sharding) and returns: SplashAttentionKernel