import torch


def biband_mask(n: int, kernel_size: int, device: torch.device, v=-1e9):
    """compute mask for local attention with kernel size.

    Args:
        n (torch.Tensor): the input length.
        kernel_size (int): The local attention kernel size.
        device (torch.device): transformer mask to the device.

    Returns: torch.Tensor. shape: [n,n]. The masked locations are -1e9
        and unmasked locations are 0.

    """
    if kernel_size is None:
        return None
    half = kernel_size // 2
    mask1 = torch.ones(n, n).triu(diagonal=-half)
    mask2 = torch.ones(n, n).tril(diagonal=half)
    mask = mask1 * mask2
    mask = (1 - mask) * v
    return mask.to(device)