code: def test_gen_attn_mask(ctx):
      class GenSelfAttnMask(HybridBlock):
          def __init__(self, dtype, layout, attn_type):
              super().__init__()
              self._dtype = dtype
              self._layout = layout
              self._attn_type = attn_type
  
          def forward(self, data, valid_length):
              return gen_self_attn_mask(
                  data,
                  valid_length,
                  dtype=self._dtype,
                  layout=self._layout,
                  attn_type=self._attn_type,
              )
  
      class GenMemAttnMask(HybridBlock):
          def __init__(self, dtype, layout):
              super().__init__()
              self._dtype = dtype
              self._layout = layout
  
          def forward(self, mem, mem_valid_length, data, valid_length):
              return gen_mem_attn_mask(
                  mem,
                  mem_valid_length,
                  data,
                  valid_length,
                  dtype=self._dtype,
                  layout=self._layout,
              )
  
      with ctx:
          batch_size = 4
          query_length = 8
          mem_length = 6
          nchannel = 5
          data = mx.np.random.normal(
              0,
              1,
              (batch_size, query_length, nchannel),
              dtype=np.float32,
          )
          valid_length = mx.np.random.randint(
              1, query_length + 1, (batch_size,)
          )
  
          mem = mx.np.random.normal(
              0, 1, (batch_size, mem_length, nchannel), dtype=np.float32
          )
          mem_valid_length = mx.np.random.randint(
              1, mem_length + 1, (batch_size,)
          )
  
          for hybridize in [False, True]:
              # Test Full Attention Mask
              mask_gen_nt = GenSelfAttnMask(
                  dtype=np.float32, layout="NT", attn_type="full"
              )
              mask_gen_tn = GenSelfAttnMask(
                  dtype=np.float32, layout="TN", attn_type="full"
              )
              if hybridize:
                  mask_gen_nt.hybridize()
                  mask_gen_tn.hybridize()
              mask_nt = mask_gen_nt(data, valid_length)
              mask_nt = mask_nt.asnumpy()
              mask_tn = mask_gen_tn(
                  mx.np.swapaxes(data, 0, 1), valid_length
              )
              mask_tn = mask_tn.asnumpy()
              mask = mask_nt
              assert_allclose(mask_nt, mask_tn)
              for b in range(batch_size):
                  v_l = valid_length.asnumpy()[b]
                  for i in range(v_l):
                      assert (mask[b, i, :v_l] == 1).all()
                      assert (mask[b, i, v_l:] == 0).all()
                  for i in range(v_l, query_length):
                      assert (mask[b, i, :] == 0).all()
  
              # Test Causal Attention Mask
              mask_gen_nt = GenSelfAttnMask(
                  dtype=np.float32, layout="NT", attn_type="causal"
              )
              mask_gen_tn = GenSelfAttnMask(
                  dtype=np.float32, layout="TN", attn_type="causal"
              )
              if hybridize:
                  mask_gen_nt.hybridize()
                  mask_gen_tn.hybridize()
              mask_nt = mask_gen_nt(data, valid_length)
              mask_tn = mask_gen_tn(
                  mx.np.swapaxes(data, 0, 1), valid_length
              )
              assert_allclose(mask_nt.asnumpy(), mask_tn.asnumpy())
              mask = mask_nt.asnumpy()
              for b in range(batch_size):
                  v_l = valid_length.asnumpy()[b]
                  for i in range(v_l):
                      assert (mask[b, i, : (i + 1)] == 1).all()
                      assert (mask[b, i, (i + 1) :] == 0).all()
                  for i in range(v_l, query_length):
                      assert (mask[b, i, :] == 0).all()
  
              # Test Mem Attention Mask
              mask_gen_nt = GenMemAttnMask(
                  dtype=np.float32, layout="NT"
              )
              mask_gen_tn = GenMemAttnMask(
                  dtype=np.float32, layout="TN"
              )
              if hybridize:
                  mask_gen_nt.hybridize()
                  mask_gen_tn.hybridize()
              mask_nt = mask_gen_nt(
                  mem, mem_valid_length, data, valid_length
              )
              mask_tn = mask_gen_tn(
                  mx.np.swapaxes(mem, 0, 1),
                  mem_valid_length,
                  mx.np.swapaxes(data, 0, 1),
                  valid_length,
              )
              mask = mask_nt.asnumpy()
              assert_allclose(mask_nt.asnumpy(), mask_tn.asnumpy())
              for b in range(batch_size):
                  data_v_l = valid_length.asnumpy()[b]
                  mem_v_l = mem_valid_length.asnumpy()[b]
                  for i in range(data_v_l):
                      assert (mask[b, i, :mem_v_l] == 1).all()
                      assert (mask[b, i, mem_v_l:] == 0).all()
                  for i in range(data_v_l, query_length):
                      assert (mask[b, i, :] == 0).all()
  
file_path: gluon-nlp/tests/test_attention_cell.py
imports: import numpy
  from numpy.testing import assert_allclose
  import pytest
  import mxnet
  from mxnet.gluon import HybridBlock
  from gluonnlp.attention_cell import multi_head_dot_attn, gen_self_attn_mask, gen_mem_attn_mask, MultiHeadAttentionCell, RelAttentionScoreCell
  from gluonnlp.utils.parameter import grad_global_norm
label_desc: Use image node_id 3 for calling a global function with example usage: test_gen_attn_mask(ctx) and returns: gen_self_attn_mask, gen_mem_attn_mask