code: def test_attention_masking(model_dict):
      # make args
      model_args = make_transformer_args(**model_dict)
  
      # setup batch
      idim = 5
      odim = 10
      ilens = [10, 5]
      olens = [20, 15]
      batch = prepare_inputs(idim, odim, ilens, olens)
  
      # define model
      model = Transformer(idim, odim, Namespace(**model_args))
  
      # test encoder self-attention
      xs = model.encoder.embed(batch["xs"])
      xs[1, ilens[1] :] = float("nan")
      x_masks = model._source_mask(batch["ilens"])
      a = model.encoder.encoders[0].self_attn
      a(xs, xs, xs, x_masks)
      aws = a.attn.detach().numpy()
      for aw, ilen in zip(aws, batch["ilens"]):
          assert not np.isnan(aw[:, :ilen, :ilen]).any()
          np.testing.assert_almost_equal(
              aw[:, :ilen, :ilen].sum(),
              float(aw.shape[0] * ilen),
              decimal=4,
          )
          assert aw[:, ilen:, ilen:].sum() == 0.0
  
      # test encoder-decoder attention
      ys = model.decoder.embed(batch["ys"])
      ys[1, olens[1] :] = float("nan")
      xy_masks = x_masks
      a = model.decoder.decoders[0].src_attn
      a(ys, xs, xs, xy_masks)
      aws = a.attn.detach().numpy()
      for aw, ilen, olen in zip(aws, batch["ilens"], batch["olens"]):
          assert not np.isnan(aw[:, :olen, :ilen]).any()
          np.testing.assert_almost_equal(
              aw[:, :olen, :ilen].sum(),
              float(aw.shape[0] * olen),
              decimal=4,
          )
          assert aw[:, olen:, ilen:].sum() == 0.0
  
      # test decoder self-attention
      y_masks = model._target_mask(batch["olens"])
      a = model.decoder.decoders[0].self_attn
      a(ys, ys, ys, y_masks)
      aws = a.attn.detach().numpy()
      for aw, olen in zip(aws, batch["olens"]):
          assert not np.isnan(aw[:, :olen, :olen]).any()
          np.testing.assert_almost_equal(
              aw[:, :olen, :olen].sum(),
              float(aw.shape[0] * olen),
              decimal=4,
          )
          assert aw[:, olen:, olen:].sum() == 0.0
  
file_path: espnet/test/test_e2e_tts_transformer.py
imports: from argparse import Namespace
  import numpy
  import pytest
  import torch
  from espnet.nets.pytorch_backend.e2e_tts_transformer import Transformer, subsequent_mask
  from espnet.nets.pytorch_backend.nets_utils import pad_list
label_desc: Use image node_id 7 for calling a global function with example usage: test_attention_masking(model_dict) without return types