import torch


def sample_input_mask(batch_size, context_length):
    """Samples a binary mask for the inputs of size batch_size x context_len
    For each row, at most one element will be 1.
    """
    mask_idx = torch.randint(context_length, size=(batch_size,))
    mask = torch.zeros((batch_size, context_length),
                       dtype=torch.int32)  # Convert to one hot B x N, B batch size, N context len
    mask[torch.arange(batch_size), mask_idx] = 1
    return mask