class_docstr: Check the construction of MaskInfo from Mask.
class_docstr_tok: ["Check","the","construction","of","MaskInfo","from","Mask","."]
code: def test_full_mask(self, is_lazy_mask: bool):
      sequence_lengths = (64, 64)
      block_shape = (16, 16)
  
      if is_lazy_mask:
          full_mask = mask_lib.MultiHeadMask(
              (mask_lib.FullMask(sequence_lengths),)
          )
      else:
          full_mask = mask_lib.MultiHeadMask(
              (
                  mask_lib.NumpyMask(
                      np.ones(sequence_lengths, dtype=np.bool_)
                  ),
              )
          )
  
      mask_info, mask_info_dkv, mask_function = self._process_mask(
          full_mask, block_shape
      )
      self.assertIsNone(mask_function)
  
      expected_mask_info = mask_info_lib.MaskInfo(
          self._expected_full_data_next[None],
          None,
          self._expected_full_block_mask[None],
          None,
          None,
      )
  
      expected_mask_info_dkv = mask_info_lib.MaskInfo(
          self._expected_full_data_next_dkv[None],
          None,
          self._expected_full_block_mask[None],
          None,
          None,
      )
  
      self._assert_mask_info_match(mask_info, expected_mask_info)
      self._assert_mask_info_match(
          mask_info_dkv, expected_mask_info_dkv
      )
  
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 7 for calling the SplashAttentionMaskInfoTest obj's underlying member method code with example usage: obj.test_full_mask(is_lazy_mask) without return types