import torch


def reduce_ctxt(
    ctxt: torch.Tensor, mask: torch.Tensor, reduction_type: str
) -> torch.Tensor:
    """
    Reduce ctxt tensor to appropriate dimension for bi-encoder scoring.

    Functionality copied from TransformerEncoder.reduce_output.

    :param ctxt:
        all transformer outputs.
    :param mask:
        context mask
    :param reduction_type:
        how to reduce the context

    :return reduced:
        return reduced context tensor
    """
    if reduction_type == 'first':
        return ctxt[:, 0, :]
    elif reduction_type == 'max':
        return ctxt.max(dim=1)[0]
    elif reduction_type == 'mean':
        divisor = mask.float().sum(dim=1).unsqueeze(-1).clamp(min=1).type_as(ctxt)
        output = ctxt.sum(dim=1) / divisor
        return output
    else:
        raise ValueError("Can't handle --reduction-type {}".format(reduction_type))