import torch


def compute_bool_rel_mat(rel_mat, vs, ignore_zero=True):
    """
    Compute boolean matrix of rel_mat. The function is called in
    `compute_rel_attn_value`.

    :param rel_mat: (bs, Lq, Lk)
    :param vs: vocabulary size, int
    :param ignore_zero: boolean. Whether to treat 0 as a padding relation.
    :return: (b, Lq, vs, Lk) or (b, Lq, vs - 1, Lk) if ignore_zero is True.
        float tensor
    """
    start_ind = 1 if ignore_zero else 0
    return torch.stack([rel_mat == i
                        for i in range(start_ind, vs)], -2).float()