import torch


def generate_sparse_one_hot(num_ents, dtype=torch.float32):
    """ Creates a two-dimensional sparse tensor with ones along the diagnoal as one-hot encoding. """
    diag_size = num_ents
    diag_range = list(range(num_ents))
    diag_range = torch.tensor(diag_range)

    return torch.sparse_coo_tensor(
        indices=torch.vstack([diag_range, diag_range]),
        values=torch.ones(diag_size, dtype=dtype),
        size=(diag_size, diag_size))