code: def _slice_mask_info(
      *,
      block_mask: np.ndarray,
      data_next: np.ndarray,
      mask_next: np.ndarray,
      head_shards: int,
      slice_function: Callable[[np.ndarray], np.ndarray],
  ):
      new_block_mask = []
      new_data_next = []
      new_mask_next = []
      for head_shard in range(head_shards):
          head_block_mask = block_mask[head_shard]
          head_block_mask = slice_function(head_block_mask)
          new_block_mask.append(head_block_mask)
  
          head_data_next = data_next[head_shard]
          head_data_next = slice_function(head_data_next)
          new_data_next.append(head_data_next)
  
          if mask_next is not None:
              head_mask_next = mask_next[head_shard]
              head_mask_next = slice_function(head_mask_next)
              new_mask_next.append(head_mask_next)
  
      block_mask = np.stack(new_block_mask, axis=0)
      data_next = np.stack(new_data_next, axis=0)
      if mask_next is not None:
          mask_next = np.stack(new_mask_next, axis=0)
  
      return block_mask, data_next, mask_next
  
file_path: jax/jax/experimental/pallas/ops/tpu/splash_attention/splash_attention_mask_info.py
imports: from __future__ import annotations
  import collections
  import functools
  from typing import Callable, Dict, List, NamedTuple, Set, Tuple
  from jax import util
  from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask
  import numpy
label_desc: Use image node_id 7 for calling a global function with example usage: _slice_mask_info() and returns: block_mask,  data_next,  mask_next