code: def _make_pad_mask_traceable(lengths, xs, length_dim, maxlen=None):
      """
      Make mask tensor containing indices of padded part.
      This is a simplified implementation of make_pad_mask without the xs input
      that supports JIT tracing for applications like exporting models to ONNX.
      Dimension length of xs should be 2 or 3
      This function will create torch.ones(maxlen, maxlen).triu(diagonal=1) and
      select rows to create mask tensor.
      """
  
      if xs is None:
          device = lengths.device
      else:
          device = xs.device
  
      if xs is not None and len(xs.shape) == 3:
          if length_dim == 1:
              lengths = lengths.unsqueeze(1).expand(
                  *xs.transpose(1, 2).shape[:2]
              )
          else:
              # Then length_dim is 2 or -1.
              if length_dim not in (-1, 2):
                  logging.warn(
                      f"Invalid length_dim {length_dim}."
                      + "We set it to -1, which is the default value."
                  )
                  length_dim = -1
              lengths = lengths.unsqueeze(1).expand(*xs.shape[:2])
  
      if maxlen is not None:
          assert xs is None
          assert maxlen >= lengths.max()
      elif xs is not None:
          maxlen = xs.shape[length_dim]
      else:
          maxlen = lengths.max()
  
      # clip max(length) to maxlen
      lengths = torch.clamp(lengths, max=maxlen).type(torch.long)
  
      mask = torch.ones(
          maxlen + 1, maxlen + 1, dtype=torch.bool, device=device
      )
      mask = triu_onnx(mask)[
          1:, :-1
      ]  # onnx cannot handle diagonal argument.
      mask = mask[lengths - 1][..., :maxlen]
  
      if xs is not None and len(xs.shape) == 3 and length_dim == 1:
          return mask.transpose(1, 2)
      else:
          return mask
  
file_path: espnet/espnet/nets/pytorch_backend/nets_utils.py
imports: import logging
  from typing import Dict
  import numpy
  import torch
label_desc: Use image node_id 5 for calling a global function with example usage: _make_pad_mask_traceable(lengths, xs, length_dim, maxlen) and returns: mask, mask