import torch


def get_bins_loss(gt_bins, pred_bins, loss_type):
    """
    gt_bins: [B, K+1]
    pred_bins: [B, K+1]

    compute cross entropy loss between gt and pred bins
    Loss: Cross entropy (CE), or L1, or L2
    """
    if loss_type == "CE":
        loss = -1 * torch.sum(gt_bins * torch.log(pred_bins + 1e-8), dim=-1)  # [B]
    elif loss_type == "L1":
        loss = torch.mean(torch.abs(gt_bins - pred_bins), dim=-1)  # [B]
    elif loss_type == "L2":
        loss = torch.mean((gt_bins - pred_bins) ** 2, dim=-1)  # [B]
    else:
        assert False, "Unknown loss type"

    return loss