code: def test_local_attention_mask(self, make_local_attention_mask):
      expected = np.array([[1]], dtype=np.bool_)
      actual = make_local_attention_mask((1, 1), (0, None), offset=0)
      self.assertArraysEqual(actual, expected)
  
      expected = np.array(
          [
              [1, 1, 1, 1],
              [1, 1, 1, 1],
              [0, 1, 1, 1],
              [0, 0, 1, 1],
          ],
          dtype=np.bool_,
      )
      actual = make_local_attention_mask((4, 4), (1, None), offset=0)
      with self.subTest("left_1"):
          self.assertArraysEqual(actual, expected)
          self.assertArraysEqual(actual, expected)
  
      expected = np.array(
          [
              [1, 1, 1, 0],
              [1, 1, 1, 1],
              [1, 1, 1, 1],
              [1, 1, 1, 1],
          ],
          dtype=np.bool_,
      )
      actual = make_local_attention_mask((4, 4), (None, 2), offset=0)
      with self.subTest("right_2"):
          self.assertArraysEqual(actual, expected)
          self.assertArraysEqual(actual, expected)
  
      expected = np.array(
          [
              [1, 1, 0, 0],
              [1, 1, 1, 0],
              [0, 1, 1, 1],
              [0, 0, 1, 1],
          ],
          dtype=np.bool_,
      )
      actual = make_local_attention_mask((4, 4), (1, 1), offset=0)
      with self.subTest("left_1_right_1"):
          self.assertArraysEqual(actual, expected)
          self.assertArraysEqual(actual, expected)
  
      expected = np.array(
          [
              [1, 0, 0, 0],
              [1, 1, 0, 0],
              [0, 1, 1, 0],
              [0, 0, 1, 1],
          ],
          dtype=np.bool_,
      )
      actual = make_local_attention_mask((4, 4), (1, 0), offset=0)
      with self.subTest("left_1_right_0"):
          self.assertArraysEqual(actual, expected)
          self.assertArraysEqual(actual, expected)
  
      expected = np.array(
          [
              [1, 1, 1, 0],
              [0, 1, 1, 1],
              [0, 0, 1, 1],
              [0, 0, 0, 1],
          ],
          dtype=np.bool_,
      )
      actual = make_local_attention_mask((4, 4), (0, 2), offset=0)
      with self.subTest("left_0_right_2"):
          self.assertArraysEqual(actual, expected)
  
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 2 for calling the SplashAttentionMaskTest obj's underlying member method code with example usage: obj.test_local_attention_mask(make_local_attention_mask) without return types