import torch


def idxes_to_onehot(idxes: torch.Tensor, vocab_size: int) -> torch.Tensor:
    """
    Given incoming tensor with dimensions [*] and vocab_size vocab_size,
    adds additional final dimension of size vocab_size, and converts each
    index into one-hot over this dimension
    """
    onehot_size = list(idxes.size()) + [vocab_size]
    onehot = torch.zeros(onehot_size, device=idxes.device, dtype=torch.float32)
    onehot.scatter_(dim=-1, index=idxes.unsqueeze(-1), value=1.0)
    return onehot