code: def test_lazy_multi_head_mask(self, shape: tuple[int, int]):
      mask_1 = mask_lib.make_random_mask(shape, 0.5, seed=1)
      mask_2 = mask_lib.make_random_mask(shape, 0.5, seed=2)
  
      lazy_multi_head = mask_lib.MultiHeadMask(
          (mask_lib.NumpyMask(mask_1), mask_lib.NumpyMask(mask_2))
      )
      dense = np.stack((mask_1, mask_2), axis=0)
  
      self._compare_masks(dense, lazy_multi_head, (256, 256))
  
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 10 for calling the SplashAttentionMaskTest obj's underlying member method code with example usage: obj.test_lazy_multi_head_mask(shape) without return types