import torch


def torch_find(index, query, vocab_size):
    """
    Finds elements of query from index, outputting the last (max) index for each
    query.
    preconditions:  (1) index and query are flat arrays (can be different sizes)
                    (2) all tokens in index and query have values < vocab_size
    """
    full_to_index = (torch.zeros(vocab_size)).long()
    if torch.cuda.is_available():
        full_to_index = full_to_index.cuda()
    index_shape_range = torch.arange(index.shape[0]).long()
    if torch.cuda.is_available():
        index_shape_range = index_shape_range.cuda()
    full_to_index[index] = index_shape_range
    result = full_to_index[query]
    return result