import torch


def depth_softargmin(cost, gamma):
    """Return a depth map with shape [N, 1, H, W].

    Arguments:
        cost: 3D cost volume tensor of shape [N, D, H, W]
        gamma: discrete depth values [D]

    Returns:
        a depth map with shape [N, 1, H, W]
    """
    return torch.sum(cost * gamma[:, None, None], 1, keepdim=True)