class_docstr: Check the construction of MaskInfo from Mask.
class_docstr_tok: ["Check","the","construction","of","MaskInfo","from","Mask","."]
code: def test_two_head_shards_two_causal_two_local(
      self, is_lazy_mask: bool
  ):
      sequence_lengths = (64, 64)
      block_shape = (16, 16)
      window_size = 8
  
      if is_lazy_mask:
          causal_mask = mask_lib.CausalMask(sequence_lengths)
          local_mask = mask_lib.LocalMask(
              sequence_lengths,
              window_size=(window_size, window_size),
              offset=0,
          )
      else:
          causal_mask = mask_lib.NumpyMask(
              mask_lib.make_causal_mask(sequence_lengths)
          )
          local_mask = mask_lib.NumpyMask(
              mask_lib.make_local_attention_mask(
                  sequence_lengths,
                  window_size=(window_size, window_size),
                  offset=0,
              )
          )
  
      mask = mask_lib.MultiHeadMask(
          (causal_mask, causal_mask, local_mask, local_mask)
      )
  
      mask_info, mask_info_dkv, mask_function = self._process_mask(
          mask, block_shape, head_shards=2
      )
      self.assertIsNone(mask_function)
  
      expected_block_mask = self._stack(
          [
              self._expected_causal_block_mask,
              self._expected_local_block_mask,
          ]
      )
  
      expected_data_next = self._stack(
          [
              self._expected_causal_data_next,
              self._expected_local_data_next,
          ]
      )
  
      expected_mask_next = self._stack(
          [
              self._expected_causal_mask_next(0),
              self._expected_local_mask_next(1),
          ]
      )
  
      expected_partial_mask_blocks = self._stack(
          [
              np.tril(np.ones(block_shape, dtype=np.int32)),
              np.triu(
                  np.tri(*block_shape, window_size, dtype=np.int32),
                  -window_size,
              ),
              np.tri(*block_shape, -window_size, dtype=np.int32),
              np.triu(
                  np.ones(block_shape, dtype=np.int32), window_size
              ),
          ],
      )
  
      expected_mask_info = mask_info_lib.MaskInfo(
          expected_data_next,
          expected_mask_next,
          expected_block_mask,
          expected_partial_mask_blocks,
          None,
      )
  
      expected_block_mask_dkv = self._stack(
          [
              self._expected_causal_block_mask_dkv,
              self._expected_local_block_mask_dkv,
          ]
      )
  
      expected_data_next_dkv = self._stack(
          [
              self._expected_causal_data_next_dkv,
              self._expected_local_data_next_dkv,
          ]
      )
  
      expected_mask_next_dkv = self._stack(
          [
              self._expected_causal_mask_next_dkv(0),
              self._expected_local_mask_next_dkv(1),
          ]
      )
  
      expected_mask_info_dkv = mask_info_lib.MaskInfo(
          expected_data_next_dkv,
          expected_mask_next_dkv,
          expected_block_mask_dkv,
          expected_partial_mask_blocks.swapaxes(-1, -2),
          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 19 for calling the SplashAttentionMaskInfoTest obj's underlying member method code with example usage: obj.test_two_head_shards_two_causal_two_local(is_lazy_mask) without return types