import torch


def _calc_dices(input, target, eps=0.001, keepdim=False):
    """Calculate dices for each sample and each channel

    Args:
        input (torch.FloatTensor): The input tensor
        target (torch.FloatTensor): The target tensor, one_hot encoded

    Returns:
        dices (torch.FloatTensor): The dices of each sample (first dim) and each
            channel (second dim)

    """
    spatial_dims = tuple(range(2 - len(input.shape), 0))
    intersection = torch.sum(input * target, dim=spatial_dims, keepdim=keepdim)
    sum1 = torch.sum(input, dim=spatial_dims, keepdim=keepdim)
    sum2 = torch.sum(target, dim=spatial_dims, keepdim=keepdim)
    dices = (2 * intersection + eps) / (sum1 + sum2 + eps)
    return dices