class_docstr: Check the construction of MaskInfo from Mask.
class_docstr_tok: ["Check","the","construction","of","MaskInfo","from","Mask","."]
code: def test_two_causal_masks(self, is_lazy_mask: bool):
      sequence_lengths = (64, 64)
      block_shape = (16, 16)
  
      if is_lazy_mask:
          causal_mask = mask_lib.CausalMask(sequence_lengths)
      else:
          causal_mask = mask_lib.NumpyMask(
              mask_lib.make_causal_mask(sequence_lengths)
          )
  
      multi_head = mask_lib.MultiHeadMask((causal_mask, causal_mask))
  
      mask_info, mask_info_dkv, mask_function = self._process_mask(
          multi_head, block_shape
      )
      if is_lazy_mask:
          self.assertIsNotNone(mask_function)
      else:
          self.assertIsNone(mask_function)
  
      expected_mask_info = mask_info_lib.MaskInfo(
          self._expected_causal_data_next[None],
          self._expected_causal_mask_next(0)[None]
          if not is_lazy_mask
          else None,
          self._expected_causal_block_mask[None],
          np.expand_dims(
              np.tril(np.ones(block_shape, dtype=np.int32)), 0
          )
          if not is_lazy_mask
          else None,
          np.arange(sequence_lengths[0], dtype=np.int32)
          if is_lazy_mask
          else None,
      )
  
      expected_mask_info_dkv = mask_info_lib.MaskInfo(
          self._expected_causal_data_next_dkv[None],
          self._expected_causal_mask_next_dkv(0)[None]
          if not is_lazy_mask
          else None,
          self._expected_causal_block_mask_dkv[None],
          np.expand_dims(
              np.tril(np.ones(block_shape, dtype=np.int32)), 0
          ).swapaxes(-1, -2)
          if not is_lazy_mask
          else None,
          np.arange(sequence_lengths[0], dtype=np.int32)
          if is_lazy_mask
          else 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 8 for calling the SplashAttentionMaskInfoTest obj's underlying member method code with example usage: obj.test_two_causal_masks(is_lazy_mask) without return types