import torch


def ones_mask(x):
    """ Creates a mask of ones in the same batch and length size as input. """
    assert len(x.size()) == 3
    return torch.ones((x.size(0), 1, x.size(-1)), dtype=x.dtype, device=x.device)