def masked_cost(cost, mask):
    """
    Mask the values from a given cost.

    Parameters
    ----------
    cost : tensor, shape 2D or 3D
        The original cost out of some cost function

    mask : tensor, shape 2D or 3D
        The mask to use to cancel part of the cost

    Returns
    -------
    masked_cost : tensor, shape cost.shape
        The masked cost

    """
    return cost * mask