def ndcg(
    outputs: torch.Tensor,
    targets: torch.Tensor,
    topk: List[int],
    gain_function="exp_rank",
) -> List[torch.Tensor]:
    """
    Computes nDCG@topk for the specified values of `topk`.

    Args:
        outputs (torch.Tensor): model outputs, logits
            with shape [batch_size; slate_size]
        targets (torch.Tensor): ground truth, labels
            with shape [batch_size; slate_size]
        gain_function:
            callable, gain function for the ground truth labels.
            Two options available:
            - `exp_rank`: torch.pow(2, x) - 1
            - `linear_rank`: x
            On the default, `exp_rank` is used
            to emphasize on retrieving the relevant documents.
        topk (List[int]):
            Parameter fro evaluation on top-k items

    Returns:
        results (Tuple[float]):
            tuple with computed ndcg@topk

    Examples:

    .. code-block:: python

        import torch
        from catalyst import metrics
        metrics.ndcg(
            outputs = torch.tensor([
                [0.5, 0.2, 0.1],
                [0.5, 0.2, 0.1],
            ]),
            targets = torch.Tensor([
                [1.0, 0.0, 1.0],
                [1.0, 0.0, 1.0],
            ]),
            topk=[2],
            gain_function="exp_rank",
        )
        # [tensor(0.6131)]

    .. code-block:: python

        import torch
        from catalyst import metrics
        metrics.ndcg(
            outputs = torch.tensor([
                [0.5, 0.2, 0.1],
                [0.5, 0.2, 0.1],
            ]),
            targets = torch.Tensor([
                [1.0, 0.0, 1.0],
                [1.0, 0.0, 1.0],
            ]),
            topk=[2],
            gain_function="exp_rank",
        )
        # [tensor(0.5000)]
    """
    results = []
    for k in topk:
        ideal_dcgs = dcg(targets, targets, gain_function)[:, :k]
        predicted_dcgs = dcg(outputs, targets, gain_function)[:, :k]
        ideal_dcgs_score = torch.sum(ideal_dcgs, dim=1)
        predicted_dcgs_score = torch.sum(predicted_dcgs, dim=1)
        ndcg_score = predicted_dcgs_score / ideal_dcgs_score
        idcg_mask = ideal_dcgs_score == 0
        ndcg_score[idcg_mask] = 0.0
        results.append(torch.mean(ndcg_score))
    return results

    Computes nDCG@topk for the specified values of `topk`.

    Args:
        outputs (torch.Tensor): model outputs, logits
            with shape [batch_size; slate_size]
        targets (torch.Tensor): ground truth, labels
            with shape [batch_size; slate_size]
        gain_function:
            callable, gain function for the ground truth labels.
            Two options available:
            - `exp_rank`: torch.pow(2, x) - 1
            - `linear_rank`: x
            On the default, `exp_rank` is used
            to emphasize on retrieving the relevant documents.
        topk (List[int]):
            Parameter fro evaluation on top-k items

    Returns:
        results (Tuple[float]):
            tuple with computed ndcg@topk

    Examples:

    .. code-block:: python

        import torch
        from catalyst import metrics
        metrics.ndcg(
            outputs = torch.tensor([
                [0.5, 0.2, 0.1],
                [0.5, 0.2, 0.1],
            ]),
            targets = torch.Tensor([
                [1.0, 0.0, 1.0],
                [1.0, 0.0, 1.0],
            ]),
            topk=[2],
            gain_function="exp_rank",
        )
        # [tensor(0.6131)]

    .. code-block:: python

        import torch
        from catalyst import metrics
        metrics.ndcg(
            outputs = torch.tensor([
                [0.5, 0.2, 0.1],
                [0.5, 0.2, 0.1],
            ]),
            targets = torch.Tensor([
                [1.0, 0.0, 1.0],
                [1.0, 0.0, 1.0],
            ]),
            topk=[2],
            gain_function="exp_rank",
        )
        # [tensor(0.5000)]
    
catalyst/metrics/functional/_ndcg.py