import torch


def half_masker(batch_size, mask_shape, dim=0):
    """Return a mask which masks the top half features of `dim`."""
    mask = torch.zeros(mask_shape).bool()
    slcs = [slice(None)] * (len(mask_shape))
    slcs[dim] = slice(0, mask_shape[dim] // 2)
    mask[slcs] = 1
    # share memory
    return mask.unsqueeze(-1).expand(batch_size, *mask_shape, 1)