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

    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.
    
train/target_speaker_extraction/models/av_tfgridnetV3/espnet2/nets_utils.py