code: def _compare_masks(
      self,
      dense_mask: np.ndarray,
      lazy_mask: mask_lib.Mask,
      block_size: tuple[int, int],
  ):
      self.assertEqual(dense_mask.shape, lazy_mask.shape)
  
      *prefix, width, height = dense_mask.shape
  
      assert width % block_size[0] == 0
      assert height % block_size[1] == 0
  
      full_lazy_mask = lazy_mask[
          (*[slice(p) for p in prefix], slice(None), slice(None))
      ]
      self.assertArraysEqual(dense_mask, full_lazy_mask)
      for i, j in np.ndindex(
          width // block_size[0], height // block_size[1]
      ):
          indexer = (
              *[slice(p) for p in prefix],
              slice(i * block_size[0], (i + 1) * block_size[0]),
              slice(j * block_size[1], (j + 1) * block_size[1]),
          )
          dense_chunk = dense_mask[indexer]
          lazy_chunk = lazy_mask[indexer]
          self.assertArraysEqual(dense_chunk, lazy_chunk)
  
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 12 for calling the SplashAttentionMaskTest obj's underlying member method code with example usage: obj._compare_masks(dense_mask, lazy_mask, block_size) without return types