class_docstr: Check the construction of MaskInfo from Mask.
class_docstr_tok: ["Check","the","construction","of","MaskInfo","from","Mask","."]
code: def test_huge_mask2(self):
      sequence_lengths = (32 * 1024, 32 * 1024)
      block_shape = (1024, 1024)
      window_size = 8
  
      local_mask = mask_lib.LocalMask(
          sequence_lengths,
          window_size=(window_size, window_size),
          offset=0,
      )
  
      multi_head = mask_lib.MultiHeadMask((local_mask,) * 32)
  
      mask_info, mask_function = mask_info_lib.process_mask(
          multi_head, block_shape
      )
  
      self.assertIsNone(mask_function)
      self.assertIsNotNone(mask_info.block_mask)
      self.assertIsNotNone(mask_info.data_next)
      self.assertIsNotNone(mask_info.mask_next)
      self.assertIsNotNone(mask_info.partial_mask_blocks)
  
imports: from __future__ import annotations
  from typing import List
  from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  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
  from jax._src import test_util
  import numpy
label_desc: Use image node_id 21 for calling the SplashAttentionMaskInfoTest obj's underlying member method code with example usage: obj.test_huge_mask2() without return types