class_docstr: Check the construction of MaskInfo from Mask.
class_docstr_tok: ["Check","the","construction","of","MaskInfo","from","Mask","."]
code: def _assert_mask_info_match(
      self,
      actual: mask_info_lib.MaskInfo,
      expected: mask_info_lib.MaskInfo,
  ):
      def assert_array_is_positive(array: np.ndarray | None):
          if array is None:
              return
  
          is_positive = np.all(array >= 0)
          self.assertTrue(is_positive)
  
      assert_array_is_positive(actual.mask_next)
      assert_array_is_positive(actual.partial_mask_blocks)
      assert_array_is_positive(actual.block_mask)
      assert_array_is_positive(actual.data_next)
  
      self.assertEqual(
          actual.data_next is not None, expected.data_next is not None
      )
      self.assertEqual(
          actual.block_mask is not None, expected.block_mask is not None
      )
      self.assertEqual(
          actual.mask_next is not None, expected.mask_next is not None
      )
      self.assertEqual(
          actual.partial_mask_blocks is not None,
          expected.partial_mask_blocks is not None,
      )
  
      self.assertEqual(
          actual.q_sequence is not None, expected.q_sequence is not None
      )
  
      if actual.partial_mask_blocks is not None:
          self.assertArraysEqual(
              actual.partial_mask_blocks,
              expected.partial_mask_blocks,
              err_msg="partial_mask_blocks",
              verbose=True,
          )
  
      if actual.q_sequence is not None:
          self.assertArraysEqual(
              actual.q_sequence,
              expected.q_sequence,
              err_msg="q_sequence",
              verbose=True,
          )
  
      self.assertArraysEqual(
          actual.block_mask,
          expected.block_mask,
          err_msg="block_mask",
          verbose=True,
      )
  
      if actual.data_next is not None and actual.block_mask is not None:
          self.assertEqual(
              actual.data_next.shape, actual.block_mask.shape
          )
  
      if actual.block_mask is not None:
          is_non_zero_block = np.where(
              expected.block_mask > 0, True, False
          )
  
          self.assertArraysEqual(
              np.where(is_non_zero_block, actual.data_next, -1),
              expected.data_next,
              err_msg="data_next",
              verbose=True,
          )
  
      if actual.mask_next is not None:
          is_partial_block = np.where(
              expected.block_mask == 1, True, False
          )
          self.assertArraysEqual(
              np.where(is_partial_block, actual.mask_next, -1),
              expected.mask_next,
              err_msg="mask_next",
              verbose=True,
          )
  
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 1 for calling the SplashAttentionMaskInfoTest obj's underlying member method code with example usage: obj._assert_mask_info_match(actual, expected) without return types