import torch


def attention_bias(inputs, mode, inf=-1e20):
    """ A bias tensor used in attention mechanism
    :param inputs: A tensor
    :param mode: either "causal" or "masking"
    :param inf: A float, denoting an inf value

    :returns: A tensor with shape [batch, heads, queries, memories]
    """

    if mode == "causal":
        result = torch.triu(torch.ones([inputs, inputs]), diagonal=1) * inf
        return result.reshape([1, 1, inputs, inputs])
    elif mode == "masking":
        result = (1.0 - inputs.float()) * inf
        return result.unsqueeze(1).unsqueeze(1)
    else:
        raise ValueError("Unknown mode %s" % mode)