import torch


def _compute_ece(prob, bin_mean_prob):
    """Compute the expected calibration error (ECE).

    Args:
      prob: Tensor, shape (2,num_bins), containing the probabilities over the
        {incorrect,correct}x{0,1,..,num_bins-1} events.
      bin_mean_prob: Tensor, shape (1,num_bins), containing the average
        probability within each bin.

    Returns:
      ece: Tensor, scalar, the expected calibration error.
    """
    pz_given_b = prob / torch.unsqueeze(torch.sum(prob, dim=0), 0)
    prob_correct = prob[1, :] / torch.sum(prob[1, :])
    ece = torch.sum(prob_correct * torch.abs(pz_given_b[1, :] - bin_mean_prob))

    return ece