eye.py

import torch


def eye(m, dtype=None, device=None):
    """Returns a sparse matrix with ones on the diagonal and zeros elsewhere.

    Args:
        m (int): The first dimension of sparse matrix.
        dtype (`torch.dtype`, optional): The desired data type of returned
            value vector. (default is set by `torch.set_default_tensor_type()`)
        device (`torch.device`, optional): The desired device of returned
            tensors. (default is set by `torch.set_default_tensor_type()`)

    :rtype: (:class:`LongTensor`, :class:`Tensor`)
    """

    row = torch.arange(m, dtype=torch.long, device=device)
    index = torch.stack([row, row], dim=0)

    value = torch.ones(m, dtype=dtype, device=device)

    return index, value