def get_mask(input_size, window_size, inner_size, device):
    """Get the attention mask of PAM-Naive"""
    # Get the size of all layers
    all_size = []
    all_size.append(input_size)
    second_size = math.floor(input_size / window_size)
    all_size.append(second_size)
    third_size = math.floor(second_size / window_size)
    all_size.append(third_size)
    fourth_size = math.floor(third_size / window_size)
    all_size.append(fourth_size)

    seq_length = sum(all_size)
    mask = torch.zeros(seq_length, seq_length, device=device)

    # Get the intra-scale mask of each scale
    inner_window = inner_size // 2
    # The first scale
    for i in range(input_size):
        left_side = max(i - inner_window, 0)
        right_side = min(i + inner_window + 1, input_size)
        mask[i, left_side:right_side] = 1
    # The second scale
    start = input_size
    for i in range(start, start + second_size):
        left_side = max(i - inner_window, start)
        right_side = min(i + inner_window + 1, start + second_size)
        mask[i, left_side:right_side] = 1
    # The third scale
    start = input_size + second_size
    for i in range(start, start + third_size):
        left_side = max(i - inner_window, start)
        right_side = min(i + inner_window + 1, start + third_size)
        mask[i, left_side:right_side] = 1
    # The fourth scale
    start = input_size + second_size + third_size
    for i in range(start, start + fourth_size):
        left_side = max(i - inner_window, start)
        right_side = min(i + inner_window + 1, start + fourth_size)
        mask[i, left_side:right_side] = 1

    # Get the inter-scale mask
    start = input_size
    for i in range(start, start + second_size):
        left_side = (i - input_size) * window_size
        if i == (start + second_size - 1):
            right_side = start
        else:
            right_side = (i - input_size + 1) * window_size
        mask[i, left_side:right_side] = 1
        mask[left_side:right_side, i] = 1
    # The third scale
    start = input_size + second_size
    for i in range(start, start + third_size):
        left_side = input_size + (i - start) * window_size
        if i == (start + third_size - 1):
            right_side = start
        else:
            right_side = input_size + (i - start + 1) * window_size
        mask[i, left_side:right_side] = 1
        mask[left_side:right_side, i] = 1
    # The fourth scale
    start = input_size + second_size + third_size
    for i in range(start, start + fourth_size):
        left_side = input_size + second_size + (i - start) * window_size
        if i == (start + fourth_size - 1):
            right_side = start
        else:
            right_side = input_size + second_size + (i - start + 1) * window_size
        mask[i, left_side:right_side] = 1
        mask[left_side:right_side, i] = 1

    mask = (1 - mask).bool()

    return mask, all_size
Get the attention mask of PAM-Naive
PatchTST_supervised/Formers/Pyraformer/pyraformer/graph_attention.py