class_docstr: Check the construction of MaskInfo from Mask.
class_docstr_tok: ["Check","the","construction","of","MaskInfo","from","Mask","."]
code: def test_huge_mask(self):
      # Don't go too high with the mask size to avoid timeouts. Prefer covering
      # multiple cases rather one very large one. This configuration replicates
      # a realistic training shape. In particular, a large number of head shards
      # and interleaving contribute to increasing processing time.
      sequence_length = (32 * 1024, 32 * 1024)
      block_shape = (512, 1024)
  
      num_shards = 16
      causal_mask = mask_lib.CausalMask(
          sequence_length, 0, shard_count=num_shards
      )
  
      multi_head = mask_lib.MultiHeadMask((causal_mask,) * 64)
  
      mask_info, mask_function = mask_info_lib.process_mask(
          multi_head, block_shape, head_shards=8, q_seq_shards=16
      )
  
      self.assertIsNotNone(mask_function)
      self.assertIsNotNone(mask_info.block_mask)
      self.assertIsNotNone(mask_info.data_next)
      self.assertIsNone(mask_info.mask_next)
      self.assertIsNone(mask_info.partial_mask_blocks)
      self.assertIsNotNone(mask_info.q_sequence)
  
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 20 for calling the SplashAttentionMaskInfoTest obj's underlying member method code with example usage: obj.test_huge_mask() without return types